Немножко обновил Vbach

MVSepLess_Epsilon_Colab.ipynb

@@ -75,6 +75,7 @@
         "faiss-cpu==1.11\n",
         "local-attention==1.11.1\n",
         "tenacity==9.1.2\n",
+        "pyworld\n",
         "gdown\n",
         "\"\"\"\n",
         "with open(\"requirements.txt\", \"w\", encoding=\"utf-8\") as f:\n",
@@ -146,7 +147,7 @@
         "\n",
         "input_url = \"\" # @param {\"type\":\"string\",\"placeholder\":\"Ссылка на аудио/видео\"}\n",
         "output_dir = \"/content/downloaded\" # @param {\"type\":\"string\",\"placeholder\":\"Директория для сохранения скачанного аудио\"}\n",
-        "cookies_path = \"\" # @param {\"type\":\"string\",\"placeholder\":\"Путь к cookies (для успешного скачивания с ютуба)\"}\n",
+        "cookies_path = \"\" # @param {\"type\":\"string\",\"placeholder\":\"Путь к cookies (дял успешного скачивания с ютуба)\"}\n",
         "downloaded_file = dw_yt_dlp(url=input_url, output_dir=output_dir, cookie=cookies_path)\n"
       ],
       "metadata": {
@@ -472,9 +473,15 @@
         "#@markdown * Имя модели:\n",
         "voicemodel_name = \"\" # @param {\"type\":\"string\",\"placeholder\":\"Имя модели\"}\n",
         "# @markdown ---\n",
+        "# @markdown  ### Hubert\n",
+        "# @markdown * Стэк\n",
+        "stack = \"fairseq\" # @param [\"fairseq\",\"transformers\"]\n",
+        "# @markdown * Имя модели для fairseq\n",
+        "fairseq_embedder = \"hubert_base\" # @param [\"hubert_base\",\"contentvec_base\",\"korean_hubert_base\",\"chinese_hubert_base\",\"portuguese_hubert_base\",\"japanese_hubert_base\"]\n",
+        "# @markdown * Имя модели для transformers\n",
+        "transformers_embedder = \"contentvec\" # @param [\"contentvec\",\"spin\",\"spin-v2\",\"chinese-hubert-base\",\"japanese-hubert-base\",\"korean-hubert-base\"]\n",
+        "# @markdown ---\n",
         "# @markdown  ### Настройки преобразования\n",
-        "# @markdown * Альтернативный пайплайн\n",
-        "alt_pipeline = False # @param {type:\"boolean\"}\n",
         "# @markdown * Влияние индекса\n",
         "index_rate = 0 # @param {\"type\":\"slider\",\"min\":0,\"max\":1,\"step\":0.01}\n",
         "# @markdown * Стерео режим\n",
@@ -520,7 +527,6 @@
         "    f\"--index_rate {index_rate}\",\n",
         "    f\"--output_name \\\"{output_name}\\\"\",\n",
         "    \"--format_name\",\n",
-        "    \"--alt_pipeline\" if alt_pipeline == True else \"\",\n",
         "    f\"--stereo_mode {stereo_mode}\",\n",
         "    f\"--method_pitch {method_pitch}\",\n",
         "    f\"--pitch {pitch}\",\n",
@@ -529,7 +535,8 @@
         "    f\"--rms {rms}\",\n",
         "    f\"--protect {protect}\",\n",
         "    f\"--f0_min {f0_min}\",\n",
-        "    f\"--f0_max {f0_max}\"\n",
+        "    f\"--f0_max {f0_max}\",\n",
+        "    f\"--embedder_name {fairseq_embedder}\" if stack == \"fairseq\" else f\"--embedder_name {transformers_embedder} --use_transformers\"\n",
         "]\n",
         "\n",
         "!{\" \".join(cmd)}"

mvsepless/__init__.py

@@ -1202,14 +1202,17 @@ class Inverter_UI(MVSEPLESS):
                 return None
 
 class Vbach(MVSEPLESS):
-    pitch_methods = ("rmvpe+", "fcpe", "mangio-crepe")
-    hop_length_values = (8, 512)
-    index_rates_values = (0, 1)
-    filter_radius_values = (0, 7)
-    protect_values = (0, 0.5)
-    rms_values = (0, 1)
-    f0_min_values = (50, 3000)
-    f0_max_values = (300, 6000)
+    def __init__(self):
+        self.pitch_methods = ("rmvpe+", "fcpe", "mangio-crepe", "mangio-crepe-tiny", "harvest", "pm")
+        self.hop_length_values = (8, 512)
+        self.index_rates_values = (0, 1)
+        self.filter_radius_values = (0, 7)
+        self.protect_values = (0, 0.5)
+        self.rms_values = (0, 1)
+        self.f0_min_values = (50, 3000)
+        self.f0_max_values = (300, 6000)
+        self.fairseq_embedders = list(self.vbach_model_manager.huberts_fairseq_dict.keys())
+        self.transformers_embedders = list(self.vbach_model_manager.huberts_transformers_dict.keys())
 
     def UI(self):
         with gr.Tab("Инференс"):
@@ -1251,7 +1254,7 @@ class Vbach(MVSEPLESS):
                             outputs=[hop_length]
                         )
                         def show_mangio_crepe_hop_length(pitch_method):
-                            return gr.update(visible=True if pitch_method in ["mangio-crepe"] else False)
+                            return gr.update(visible=True if pitch_method in ["mangio-crepe", "mangio-crepe-tiny"] else False)
                         stereo_mode = gr.Radio(
                             choices=["mono", "left/right", "sim/dif"],
                             label="Стерео режим",
@@ -1260,6 +1263,19 @@ class Vbach(MVSEPLESS):
                             interactive=True
                         )
                         alt_pl = gr.Checkbox(label="Альтернативный пайплайн", info="Аудио нарезается на фиксированные чанки с перекрытием, что исключает любые щелчки на выходе (исключение - если есть щелчки в самой модели из-за грязного датасета)\nРазмер чанка вычисляется на основе 40% свободной видеопамяти", value=False, interactive=True)
+
+                    with gr.Group():
+                        embedder_name = gr.Radio(label="Модель Hubert", choices=self.fairseq_embedders, value=self.fairseq_embedders[0])
+                        transformers_mode = gr.Checkbox(label="Использовать стек Transformers", value=False, interactive=True)
+                        @transformers_mode.change(
+                            inputs=[transformers_mode], outputs=[embedder_name]
+                        )
+                        def change_embedders(tr_m):
+                            if tr_m:
+                                return gr.update(value=self.transformers_embedders[0], choices=self.transformers_embedders)
+                            else:
+                                return gr.update(choices=self.fairseq_embedders, value=self.fairseq_embedders[0])
+
                     with gr.Accordion(label="Дополнительные настройки",open=False):
                         with gr.Group():
                             with gr.Row():
@@ -1297,10 +1313,12 @@ class Vbach(MVSEPLESS):
                     format_output_name_check, 
                     output_format,
                     stereo_mode,
-                    alt_pl
+                    alt_pl,
+                    embedder_name,
+                    transformers_mode
                 ], outputs=[converted_state, status]
             )
-            def vbach_convert_batch(ifl, mn, pm, p, hl, ir, fr, rms, pr, f0min, f0max, on, fn, of, sm, alt_pipeline):
+            def vbach_convert_batch(ifl, mn, pm, p, hl, ir, fr, rms, pr, f0min, f0max, on, fn, of, sm, alt_pipeline, em_n, tr_m):
                 output_converted_files = []
                 progress = gr.Progress()
                 if ifl:
@@ -1308,7 +1326,7 @@ class Vbach(MVSEPLESS):
                         try:
                             print(f"Файл {i} из {len(ifl)}: {file}")
                             progress(progress=(i / len(ifl)), desc=f"Файл {i} из {len(ifl)}")
-                            out_conv = vbach_inference(input_file=file, model_name=mn, output_dir=tempfile.mkdtemp(), output_name=on, format_name=True if len(ifl) > 1 else fn, output_format=of, pitch=p, method_pitch=pm, output_bitrate=320, add_params={ "index_rate": ir,"filter_radius": fr,"protect": pr,"rms": rms,"mangio_crepe_hop_length": hl,"f0_min": f0min,"f0_max": f0max,"stereo_mode": sm }, pipeline_mode="alt" if alt_pipeline == True else "orig")
+                            out_conv = vbach_inference(input_file=file, model_name=mn, output_dir=tempfile.mkdtemp(), output_name=on, format_name=True if len(ifl) > 1 else fn, output_format=of, pitch=p, method_pitch=pm, output_bitrate=320, add_params={ "index_rate": ir,"filter_radius": fr,"protect": pr,"rms": rms,"mangio_crepe_hop_length": hl,"f0_min": f0min,"f0_max": f0max,"stereo_mode": sm }, pipeline_mode="alt" if alt_pipeline == True else "orig", embedder_name=em_n, stack="transformers" if tr_m == True else "fairseq")
                             output_converted_files.append(out_conv)
                         except Exception as e:
                             print(e)

mvsepless/model_manager.py

@@ -254,15 +254,105 @@ class VbachModelManager:
     def __init__(self):
         self.rmvpe_path = os.path.join(script_dir, "predictors", "rmvpe.pt")
         self.fcpe_path = os.path.join(script_dir, "predictors", "fcpe.pt")
-        self.hubert_path = os.path.join(script_dir, "embedders", "hubert_base.pt")
-        self.requirements = [["https://huggingface.co/Politrees/RVC_resources/resolve/main/predictors/rmvpe.pt", self.rmvpe_path], ["https://huggingface.co/Politrees/RVC_resources/resolve/main/predictors/fcpe.pt", self.fcpe_path], ["https://huggingface.co/Politrees/RVC_resources/resolve/main/embedders/hubert_base.pt", self.hubert_path]]
+        self.custom_fairseq_huberts_dir = os.path.join(script_dir, "custom_fairseq_embedders")
+        self.custom_transformers_huberts_dir = os.path.join(script_dir, "custom_transformers_embedders")
+        self.huberts_fairseq_dict = {
+            "hubert_base": {
+                "url": "https://huggingface.co/Politrees/RVC_resources/resolve/main/embedders/hubert_base.pt",
+                "local_path": os.path.join(self.custom_fairseq_huberts_dir, "hubert_base.pt")
+            },
+            "contentvec_base": {
+                "url": "https://huggingface.co/Politrees/RVC_resources/resolve/main/embedders/contentvec_base.pt",
+                "local_path": os.path.join(self.custom_fairseq_huberts_dir, "contentvec_base.pt")
+            },
+            "korean_hubert_base": {
+                "url": "https://huggingface.co/Politrees/RVC_resources/resolve/main/embedders/korean_hubert_base.pt",
+                "local_path": os.path.join(self.custom_fairseq_huberts_dir, "korean_hubert_base.pt")
+            },
+            "chinese_hubert_base": {
+                "url": "https://huggingface.co/Politrees/RVC_resources/resolve/main/embedders/chinese_hubert_base.pt",
+                "local_path": os.path.join(self.custom_fairseq_huberts_dir, "chinese_hubert_base.pt")
+            },
+            "portuguese_hubert_base": {
+                "url": "https://huggingface.co/Politrees/RVC_resources/resolve/main/embedders/portuguese_hubert_base.pt",
+                "local_path": os.path.join(self.custom_fairseq_huberts_dir, "portuguese_hubert_base.pt")
+            },
+            "japanese_hubert_base": {
+                "url": "https://huggingface.co/Politrees/RVC_resources/resolve/main/embedders/japanese_hubert_base.pt",
+                "local_path": os.path.join(self.custom_fairseq_huberts_dir, "japanese_hubert_base.pt")
+            }
+        }
+        self.huberts_transformers_dict = {
+            "contentvec": {
+                "base_dir": os.path.join(self.custom_transformers_huberts_dir, "contentvec"),
+                "url_bin": "https://huggingface.co/IAHispano/Applio/resolve/main/Resources/embedders/contentvec/pytorch_model.bin",
+                "url_json": "https://huggingface.co/IAHispano/Applio/resolve/main/Resources/embedders/contentvec/config.json",
+                "local_bin": os.path.join(self.custom_transformers_huberts_dir, "contentvec", "pytorch_model.bin"),
+                "local_json": os.path.join(self.custom_transformers_huberts_dir, "contentvec", "config.json")
+            },
+            "spin": {
+                "base_dir": os.path.join(self.custom_transformers_huberts_dir, "spin"),
+                "url_bin": "https://huggingface.co/IAHispano/Applio/resolve/main/Resources/embedders/spin/pytorch_model.bin",
+                "url_json": "https://huggingface.co/IAHispano/Applio/resolve/main/Resources/embedders/spin/config.json",
+                "local_bin": os.path.join(self.custom_transformers_huberts_dir, "spin", "pytorch_model.bin"),
+                "local_json": os.path.join(self.custom_transformers_huberts_dir, "spin", "config.json")
+            },
+            "spin-v2": {
+                "base_dir": os.path.join(self.custom_transformers_huberts_dir, "spinv2"),
+                "url_bin": "https://huggingface.co/IAHispano/Applio/resolve/main/Resources/embedders/spin-v2/pytorch_model.bin",
+                "url_json": "https://huggingface.co/IAHispano/Applio/resolve/main/Resources/embedders/spin-v2/config.json",
+                "local_bin": os.path.join(self.custom_transformers_huberts_dir, "spinv2", "pytorch_model.bin"),
+                "local_json": os.path.join(self.custom_transformers_huberts_dir, "spinv2", "config.json")
+            },
+            "chinese-hubert-base": {
+                "base_dir": os.path.join(self.custom_transformers_huberts_dir, "chinese_hubert_base"),
+                "url_bin": "https://huggingface.co/IAHispano/Applio/resolve/main/Resources/embedders/chinese_hubert_base/pytorch_model.bin",
+                "url_json": "https://huggingface.co/IAHispano/Applio/resolve/main/Resources/embedders/chinese_hubert_base/config.json",
+                "local_bin": os.path.join(self.custom_transformers_huberts_dir, "chinese_hubert_base", "pytorch_model.bin"),
+                "local_json": os.path.join(self.custom_transformers_huberts_dir, "chinese_hubert_base", "config.json")
+            },
+            "japanese-hubert-base": {
+                "base_dir": os.path.join(self.custom_transformers_huberts_dir, "japanese_hubert_base"),
+                "url_bin": "https://huggingface.co/IAHispano/Applio/resolve/main/Resources/embedders/japanese_hubert_base/pytorch_model.bin",
+                "url_json": "https://huggingface.co/IAHispano/Applio/resolve/main/Resources/embedders/japanese_hubert_base/config.json",
+                "local_bin": os.path.join(self.custom_transformers_huberts_dir, "japanese_hubert_base", "pytorch_model.bin"),
+                "local_json": os.path.join(self.custom_transformers_huberts_dir, "japanese_hubert_base", "config.json")
+            },
+            "korean-hubert-base": {
+                "base_dir": os.path.join(self.custom_transformers_huberts_dir, "korean_hubert_base"),
+                "url_bin": "https://huggingface.co/IAHispano/Applio/resolve/main/Resources/embedders/korean_hubert_base/pytorch_model.bin",
+                "url_json": "https://huggingface.co/IAHispano/Applio/resolve/main/Resources/embedders/korean_hubert_base/config.json",
+                "local_bin": os.path.join(self.custom_transformers_huberts_dir, "korean_hubert_base", "pytorch_model.bin"),
+                "local_json": os.path.join(self.custom_transformers_huberts_dir, "korean_hubert_base", "config.json")
+            }
+        }
+        self.requirements = [["https://huggingface.co/Politrees/RVC_resources/resolve/main/predictors/rmvpe.pt", self.rmvpe_path], ["https://huggingface.co/Politrees/RVC_resources/resolve/main/predictors/fcpe.pt", self.fcpe_path]]
         self.voicemodels_dir = os.path.join(script_dir, "vbach_models_cache")
         os.makedirs(self.voicemodels_dir, exist_ok=True)
         self.voicemodels_info = os.path.join(self.voicemodels_dir, "vbach_models.json")
         self.voicemodels: Dict[str, Dict[str, str]] = {}
         self.download_requirements()
+        self.check_hubert("hubert_base")
         self.check_and_load()
         pass
+
+    def check_hubert(self, embedder_name):
+        if embedder_name in self.huberts_fairseq_dict:
+            if not os.path.exists(self.huberts_fairseq_dict[embedder_name]["local_path"]):
+                dw_file(self.huberts_fairseq_dict[embedder_name]["url"], self.huberts_fairseq_dict[embedder_name]["local_path"])
+            return self.huberts_fairseq_dict[embedder_name]["local_path"]
+        else:
+            return None
+        
+    def check_hubert_transformers(self, embedder_name):
+        if embedder_name in self.huberts_transformers_dict:
+            os.makedirs(self.huberts_transformers_dict[embedder_name]["base_dir"], exist_ok=True)
+            if not os.path.exists(self.huberts_transformers_dict[embedder_name]["local_bin"]) and not os.path.exists(self.huberts_transformers_dict[embedder_name]["local_json"]):
+                dw_file(self.huberts_transformers_dict[embedder_name]["url_bin"], self.huberts_transformers_dict[embedder_name]["local_bin"])
+                dw_file(self.huberts_transformers_dict[embedder_name]["url_json"], self.huberts_transformers_dict[embedder_name]["local_json"])
+            return self.huberts_transformers_dict[embedder_name]["base_dir"]
+        else:
+            return None
     
     def write_voicemodels_info(self):
         with open(self.voicemodels_info, "w") as f:

mvsepless/vbach_infer.py

@@ -1,6 +1,7 @@
 import os
 import gc
 import torch
+from torch import nn
 import torch.nn.functional as F
 import torchcrepe
 import faiss
@@ -9,6 +10,11 @@ import math
 import numpy as np
 from scipy import signal
 import argparse
+from functools import lru_cache
+import pyworld
+import parselmouth
+from transformers import HubertModel
+from typing import Tuple, Any, Dict
 script_dir = os.path.dirname(os.path.abspath(__file__))
 
 FILTER_ORDER = 5  
@@ -44,6 +50,26 @@ namer = Namer()
 RMVPE_DIR = model_manager.rmvpe_path
 FCPE_DIR = model_manager.fcpe_path
 
+input_audio_path2wav = {}
+
+class HubertModelWithFinalProj(HubertModel):
+    def __init__(self, config):
+        super().__init__(config)
+        self.final_proj = nn.Linear(config.hidden_size, config.classifier_proj_size)
+
+@lru_cache
+def get_harvest_f0(input_audio_path, fs, f0max, f0min, frame_period):
+    audio = input_audio_path2wav[input_audio_path]
+    f0, t = pyworld.harvest(
+        audio,
+        fs=fs,
+        f0_ceil=f0max,
+        f0_floor=f0min,
+        frame_period=frame_period,
+    )
+    f0 = pyworld.stonemask(audio, f0, t, fs)
+    return f0
+
 def remove_center(input_array, samplerate, rdf=0.99999, window_size=2048, overlap=2, window_type="blackman"):
     
     left = input_array[0]
@@ -137,7 +163,7 @@ class AudioProcessor:
 
 # Класс для преобразования голоса
 class VC:
-    def __init__(self, tgt_sr, config):
+    def __init__(self, tgt_sr, config, stack="fairseq"):
         """
         Инициализация параметров для преобразования голоса.
         """
@@ -156,10 +182,11 @@ class VC:
         self.t_max = self.sample_rate * self.x_max
         self.time_step = self.window / self.sample_rate * 1000
         self.device = config.device
+        self.vc = self._vc_transformers if stack == "transformers" else self._vc
 
-    def get_f0_crepe(self, x, f0_min, f0_max, p_len, hop_length, model="full"):
+    def get_f0_mangio_crepe(self, x, f0_min, f0_max, p_len, hop_length, model="full"):
         """
-        Получает F0 с использованием модели crepe.
+        Получает F0 с использованием модели mangio-crepe.
         """
         x = x.astype(np.float32)
         x /= np.quantile(np.abs(x), 0.999)
@@ -219,12 +246,36 @@ class VC:
         """
         Получает F0 с использованием выбранного метода.
         """
-        global inputaudio_path2wav
+        global input_audio_path2wav
+        time_step = self.window / self.sample_rate * 1000
         f0_mel_min = 1127 * np.log(1 + f0_min / 700)
         f0_mel_max = 1127 * np.log(1 + f0_max / 700)
 
-        if f0_method == "mangio-crepe":
-            f0 = self.get_f0_crepe(x, f0_min, f0_max, p_len, int(hop_length))
+        if f0_method in ["mangio-crepe", "mangio-crepe-tiny"]:
+            f0 = self.get_f0_mangio_crepe(x, f0_min, f0_max, p_len, int(hop_length), "tiny" if f0_method == "mangio-crepe-tiny" else "full")
+
+        elif f0_method == "harvest":
+            input_audio_path2wav = {}
+            input_audio_path2wav[inputaudio_path] = x.astype(np.double)
+            f0 = get_harvest_f0(inputaudio_path, self.sample_rate, f0_max, f0_min, 10)
+            if filter_radius > 2:
+                f0 = signal.medfilt(f0, 3)
+        elif f0_method == "pm":
+            f0 = (
+                parselmouth.Sound(x, self.sample_rate)
+                .to_pitch_ac(
+                    time_step=time_step / 1000,
+                    voicing_threshold=0.6,
+                    pitch_floor=f0_min,
+                    pitch_ceiling=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"
+                )
 
         elif f0_method == "rmvpe+":
             params = {
@@ -274,7 +325,7 @@ class VC:
         f0_coarse = np.rint(f0_mel).astype(int)
         return f0_coarse, f0bak
 
-    def vc(
+    def _vc(
         self,
         model,
         net_g,
@@ -375,6 +426,91 @@ class VC:
         
         return audio1
 
+    def _vc_transformers(
+        self,
+        model,
+        net_g,
+        sid,
+        audio0,
+        pitch,
+        pitchf,
+        index,
+        big_npy,
+        index_rate,
+        version,
+        protect,
+    ):
+        """
+        Performs voice conversion on a given audio segment.
+
+        Args:
+            model: The feature extractor model.
+            net_g: The generative model for synthesizing speech.
+            sid: Speaker ID for the target voice.
+            audio0: The input audio segment.
+            pitch: Quantized F0 contour for pitch guidance.
+            pitchf: Original F0 contour for pitch guidance.
+            index: FAISS index for speaker embedding retrieval.
+            big_npy: Speaker embeddings stored in a NumPy array.
+            index_rate: Blending rate for speaker embedding retrieval.
+            version: Model version (Keep to support old models).
+            protect: Protection level for preserving the original pitch.
+        """
+        with torch.no_grad():
+            pitch_guidance = pitch != None and pitchf != None
+            # prepare source audio
+            feats = torch.from_numpy(audio0).float()
+            feats = feats.mean(-1) if feats.dim() == 2 else feats
+            assert feats.dim() == 1, feats.dim()
+            feats = feats.view(1, -1).to(self.device)
+            # extract features
+            feats = model(feats)["last_hidden_state"]
+            feats = (
+                model.final_proj(feats[0]).unsqueeze(0) if version == "v1" else feats
+            )
+            # make a copy for pitch guidance and protection
+            feats0 = feats.clone() if pitch_guidance else None
+            if (
+                index
+            ):  # set by parent function, only true if index is available, loaded, and index rate > 0
+                feats = self._retrieve_speaker_embeddings(
+                    feats, index, big_npy, index_rate
+                )
+            # feature upsampling
+            feats = F.interpolate(feats.permute(0, 2, 1), scale_factor=2).permute(
+                0, 2, 1
+            )
+            # adjust the length if the audio is short
+            p_len = min(audio0.shape[0] // self.window, feats.shape[1])
+            if pitch_guidance:
+                feats0 = F.interpolate(feats0.permute(0, 2, 1), scale_factor=2).permute(
+                    0, 2, 1
+                )
+                pitch, pitchf = pitch[:, :p_len], pitchf[:, :p_len]
+                # Pitch protection blending
+                if protect < 0.5:
+                    pitchff = pitchf.clone()
+                    pitchff[pitchf > 0] = 1
+                    pitchff[pitchf < 1] = protect
+                    feats = feats * pitchff.unsqueeze(-1) + feats0 * (
+                        1 - pitchff.unsqueeze(-1)
+                    )
+                    feats = feats.to(feats0.dtype)
+            else:
+                pitch, pitchf = None, None
+            p_len = torch.tensor([p_len], device=self.device).long()
+            audio1 = (
+                (net_g.infer(feats.float(), p_len, pitch, pitchf.float(), sid)[0][0, 0])
+                .data.cpu()
+                .float()
+                .numpy()
+            )
+            # clean up
+            del feats, feats0, p_len
+            if torch.cuda.is_available():
+                torch.cuda.empty_cache()
+        return audio1
+
     def pipeline(
         self,
         model,
@@ -578,8 +714,6 @@ class VC:
         f0_min=50,
         f0_max=1100,
     ):
-        import torch
-        import numpy as np
 
         device = self.device
         audio = signal.filtfilt(bh, ah, audio)
@@ -798,7 +932,17 @@ class VC:
         # Для CPU или в случае ошибки используем консервативный размер
         return min(base_chunk_size, audio_length)
 
-
+    def _retrieve_speaker_embeddings(self, feats, index, big_npy, index_rate): # для Transformers
+        npy = feats[0].cpu().numpy()
+        score, ix = index.search(npy, k=8)
+        weight = np.square(1 / score)
+        weight /= weight.sum(axis=1, keepdims=True)
+        npy = np.sum(big_npy[ix] * np.expand_dims(weight, axis=2), axis=1)
+        feats = (
+            torch.from_numpy(npy).unsqueeze(0).to(self.device) * index_rate
+            + (1 - index_rate) * feats
+        )
+        return feats
 
 def overlay_mono_on_stereo(monoaudio, stereoaudio, gain=0.5):
     if monoaudio is None or stereoaudio is None:
@@ -992,35 +1136,128 @@ def load_hubert(device, is_half, model_path):
     hubert.eval()
     return hubert
 
-# Получение голосового преобразователя
-def get_vc(device, is_half, config, model_path):
-    cpt = torch.load(model_path, map_location="cpu", weights_only=False)
-    if "config" not in cpt or "weight" not in cpt:
-        raise ValueError(
-            f"Некорректный формат для {model_path}. "
-            "Используйте голосовую модель, обученную с использованием RVC v2."
+def get_vc(
+    device: torch.device,
+    is_half: bool,
+    config: Any,
+    model_path: str,
+    stack: Any
+) -> Tuple[Dict[str, Any], str, torch.nn.Module, int, VC, int]:
+    """
+    Загружает модель RVC для голосовой конвертации.
+    
+    Args:
+        device: Устройство для загрузки модели (CPU/GPU)
+        is_half: Использовать ли половинную точность (float16)
+        config: Конфигурация модели
+        model_path: Путь к файлу модели (.pth)
+        stack: Объект стека для инициализации VC
+        
+    Returns:
+        Tuple containing:
+        - cpt: Checkpoint модели
+        - version: Версия модели
+        - net_g: Сетевой генератор
+        - tgt_sr: Целевая частота дискретизации
+        - vc: Объект голосовой конвертации
+        - use_f0: Использование F0 (0/1)
+        
+    Raises:
+        ValueError: Если файл модели имеет некорректный формат
+        FileNotFoundError: Если файл модели не найден
+    """
+    
+    # Проверка существования файла
+    if not os.path.isfile(model_path):
+        raise FileNotFoundError(f"Файл модели не найден: {model_path}")
+    
+    try:
+        # Загружаем состояние модели
+        cpt = torch.load(model_path, map_location="cpu", weights_only=True)
+        
+        # Проверяем структуру загруженного файла
+        required_keys = ["config", "weight"]
+        missing_keys = [key for key in required_keys if key not in cpt]
+        
+        if missing_keys:
+            raise ValueError(
+                f"Некорректный формат модели {model_path}. "
+                f"Отсутствующие ключи: {missing_keys}. "
+                "Используйте модель RVC формата."
+            )
+        
+        # Извлекаем параметры модели
+        tgt_sr = cpt["config"][-1]
+        
+        # Обновляем размерность в config на основе весов
+        emb_weight_shape = cpt["weight"]["emb_g.weight"].shape
+        cpt["config"][-3] = emb_weight_shape[0]  # Количество спикеров
+        
+        # Получаем дополнительные параметры модели
+        use_f0 = cpt.get("f0", 1)
+        version = cpt.get("version", "v1")
+        vocoder = cpt.get("vocoder", "HiFi-GAN")
+        
+        # Определяем размерность входных данных в зависимости от версии
+        text_enc_hidden_dim = 768 if version == "v2" else 256
+        
+        print(f"Загружаем модель: {os.path.basename(model_path)}")
+        print(f"Версия: {version}, F0: {use_f0}, Частота: {tgt_sr}Hz")
+        print(f"Количество спикеров: {emb_weight_shape[0]}")
+        
+        # Инициализируем синтезатор
+        net_g = Synthesizer(
+            *cpt["config"],
+            use_f0=use_f0,
+            text_enc_hidden_dim=text_enc_hidden_dim,
+            vocoder=vocoder
         )
-
-    tgt_sr = cpt["config"][-1]
-    cpt["config"][-3] = cpt["weight"]["emb_g.weight"].shape[0]
-    pitch_guidance = cpt.get("f0", 1)
-    version = cpt.get("version", "v1")
-    input_dim = 768 if version == "v2" else 256
-
-    net_g = Synthesizer(
-        *cpt["config"],
-        use_f0=pitch_guidance,
-        input_dim=input_dim,
-        is_half=is_half,
-    )
-
-    del net_g.enc_q
-    print(net_g.load_state_dict(cpt["weight"], strict=False))
-    net_g.eval().to(device)
-    net_g = net_g.half() if is_half else net_g.float()
-
-    vc = VC(tgt_sr, config)
-    return cpt, version, net_g, tgt_sr, vc
+        
+        # Удаляем ненужный слой enc_q
+        if hasattr(net_g, 'enc_q'):
+            del net_g.enc_q
+        else:
+            print("Предупреждение: слой enc_q не найден в модели")
+        
+        # Загружаем веса с проверкой
+        missing_keys, unexpected_keys = net_g.load_state_dict(
+            cpt["weight"], 
+            strict=False
+        )
+        
+        if missing_keys:
+            print(f"Предупреждение: отсутствующие ключи при загрузке модели: {missing_keys}")
+        
+        if unexpected_keys:
+            print(f"Предупреждение: неожиданные ключи при загрузке модели: {unexpected_keys}")
+        
+        # Настройка модели для inference
+        net_g.eval()
+        
+        # Перемещаем модель на нужное устройство и устанавливаем точность
+        net_g = net_g.to(device)
+        if is_half:
+            net_g = net_g.half()
+            print("Модель переведена в половинную точность (float16)")
+        else:
+            net_g = net_g.float()
+            print("Модель использует полную точность (float32)")
+        
+        # Инициализируем объект конвертера голоса
+        vc = VC(tgt_sr, config, stack)
+        
+        # Очистка памяти CPU
+        if torch.cuda.is_available():
+            torch.cuda.empty_cache()
+        
+        print(f"Модель успешно загружена на устройство: {device}")
+        
+        return cpt, version, net_g, tgt_sr, vc, use_f0
+        
+    except torch.serialization.pickle.UnpicklingError as e:
+        raise ValueError(f"Файл {model_path} поврежден или имеет неверный формат") from e
+    except Exception as e:
+        raise RuntimeError(f"Ошибка при загрузке модели: {str(e)}") from e
 
 def rvc_infer(
     index_path,
@@ -1039,6 +1276,7 @@ def rvc_infer(
     hop_length,
     vc,
     hubert_model,
+    pitch_guidance,
     f0_min=50,
     f0_max=1100,
     format_output="wav",
@@ -1053,7 +1291,6 @@ def rvc_infer(
         pipeline = vc.pipeline
 
     mid, left, right = loadaudio(input_path, 16000, stereo_mode)
-    pitch_guidance = cpt.get("f0", 1)
     
     if stereo_mode == "mono":
         if mid is None:
@@ -1263,15 +1500,87 @@ def voice_conversion(
     format_output,
     output_bitrate,
     stereo_mode,
-    hubert_path=None,
+    embedder_name="hubert_base",
+    pipeline_mode="orig"
+):
+    rvc_model_path, rvc_index_path = load_rvc_model(voice_model)
+
+    config = Config()
+    hubert_path = model_manager.check_hubert(embedder_name)
+    if not hubert_path:
+        raise ValueError(
+            f"Эмбеддера {embedder_name} не существует. "
+            "Возможно, вы неправильно ввели имя."
+        )
+    hubert_model = load_hubert(config.device, config.is_half, hubert_path)
+    cpt, version, net_g, tgt_sr, vc, use_f0 = get_vc(
+        config.device, config.is_half, config, rvc_model_path, "fairseq"
+    )
+
+    outputaudio = rvc_infer(
+        rvc_index_path,
+        index_rate,
+        vocals_path,
+        output_path,
+        pitch,
+        f0_method,
+        cpt,
+        version,
+        net_g,
+        filter_radius,
+        tgt_sr,
+        volume_envelope,
+        protect,
+        hop_length,
+        vc,
+        hubert_model,
+        use_f0,
+        f0_min,
+        f0_max,
+        format_output,
+        output_bitrate,
+        stereo_mode,
+        pipeline_mode
+    )
+
+    del hubert_model, cpt, net_g, vc
+    gc.collect()
+    torch.cuda.empty_cache()
+    return outputaudio
+
+def voice_conversion_transformers(
+    voice_model,
+    vocals_path,
+    output_path,
+    pitch,
+    f0_method,
+    index_rate,
+    filter_radius,
+    volume_envelope,
+    protect,
+    hop_length,
+    f0_min,
+    f0_max,
+    format_output,
+    output_bitrate,
+    stereo_mode,
+    embedder_name="contentvec",
     pipeline_mode="orig"
 ):
     rvc_model_path, rvc_index_path = load_rvc_model(voice_model)
 
     config = Config()
-    hubert_model = load_hubert(config.device, config.is_half, hubert_path if hubert_path else model_manager.hubert_path)
-    cpt, version, net_g, tgt_sr, vc = get_vc(
-        config.device, config.is_half, config, rvc_model_path
+    hubert_path = model_manager.check_hubert_transformers(embedder_name)
+    if not hubert_path:
+        raise ValueError(
+            f"Эмбеддера {embedder_name} не существует. "
+            "Возможно, вы неправильно ввели имя."
+        )
+    #hubert_model = load_hubert(config.device, config.is_half, hubert_path)
+    hubert_model = HubertModelWithFinalProj.from_pretrained(hubert_path)
+    hubert_model = hubert_model.to(config.device)
+    cpt, version, net_g, tgt_sr, vc, use_f0 = get_vc(
+        config.device, config.is_half, config, rvc_model_path, "transformers"
     )
 
     outputaudio = rvc_infer(
@@ -1291,6 +1600,7 @@ def voice_conversion(
         hop_length,
         vc,
         hubert_model,
+        use_f0,
         f0_min,
         f0_max,
         format_output,
@@ -1304,6 +1614,7 @@ def voice_conversion(
     torch.cuda.empty_cache()
     return outputaudio
 
+
 def vbach_inference(
         input_file: str, 
         model_name: str, 
@@ -1315,6 +1626,8 @@ def vbach_inference(
         method_pitch: str, 
         format_name: bool = False,
         pipeline_mode: str = "orig",
+        embedder_name: str | None = "hubert_base",
+        stack: str = "fairseq",
         add_params: dict = {
             "index_rate": 0, 
             "filter_radius": 3, 
@@ -1326,6 +1639,13 @@ def vbach_inference(
             "stereo_mode": "mono"
         }
     ):
+
+    if stack == "fairseq":
+        vbach_convert = voice_conversion
+    elif stack == "transformers":
+        vbach_convert = voice_conversion_transformers
+
+
     stereo_mode = add_params.get("stereo_mode", "mono")
     index_rate = add_params.get("index_rate", 0)
     filter_radius = add_params.get("filter_radius", 3)
@@ -1355,7 +1675,7 @@ def vbach_inference(
         final_output_name = output_name
 
     final_output_path = os.path.join(output_dir, f"{final_output_name}.{output_format}")
-    output_converted_voice = voice_conversion(voice_model=model_name, vocals_path=input_file, output_path=final_output_path, pitch=pitch, f0_method=method_pitch, index_rate=index_rate, filter_radius=filter_radius, volume_envelope=rms, protect=protect, hop_length=mangio_crepe_hop_length, f0_min=f0_min, f0_max=f0_max, format_output=output_format, hubert_path=None, output_bitrate=output_bitrate, stereo_mode=stereo_mode, pipeline_mode=pipeline_mode)
+    output_converted_voice = vbach_convert(voice_model=model_name, vocals_path=input_file, output_path=final_output_path, pitch=pitch, f0_method=method_pitch, index_rate=index_rate, filter_radius=filter_radius, volume_envelope=rms, protect=protect, hop_length=mangio_crepe_hop_length, f0_min=f0_min, f0_max=f0_max, format_output=output_format, output_bitrate=output_bitrate, stereo_mode=stereo_mode, pipeline_mode=pipeline_mode, embedder_name=embedder_name)
     print(f"Инференс завершен\nПуть к выходному файлу: \"{output_converted_voice}\"")
     return output_converted_voice
 
@@ -1463,13 +1783,24 @@ if __name__ == "__main__":
         action="store_true",
         help="Альтернативный пайплайн",
     )
+    parser.add_argument(
+        "--use_transformers",
+        action="store_true",
+        help="Использовать transformers",
+    )
+    parser.add_argument(
+        "--embedder_name",
+        type=str,
+        default="hubert_base",
+        help="Имя Hubert модели",
+    )
 
     args = parser.parse_args()
 
     if args.input:
         if os.path.exists(args.input) and os.path.isfile(args.input):
             if audio.check(args.input):
-                vbach_inference(input_file=args.input, model_name=args.model_name, output_dir=args.output_dir, output_name=args.output_name, output_bitrate=args.output_bitrate, output_format=args.output_format, pitch=args.pitch, method_pitch=args.method_pitch, format_name=args.format_name, add_params={ "index_rate": args.index_rate,"filter_radius": args.filter_radius,"protect": args.protect,"rms": args.rms,"mangio_crepe_hop_length": args.hop_length,"f0_min": args.f0_min,"f0_max": args.f0_max,"stereo_mode": args.stereo_mode}, pipeline_mode="alt" if args.alt_pipeline == True else "orig")
+                vbach_inference(input_file=args.input, model_name=args.model_name, output_dir=args.output_dir, output_name=args.output_name, output_bitrate=args.output_bitrate, output_format=args.output_format, pitch=args.pitch, method_pitch=args.method_pitch, format_name=args.format_name, add_params={ "index_rate": args.index_rate,"filter_radius": args.filter_radius,"protect": args.protect,"rms": args.rms,"mangio_crepe_hop_length": args.hop_length,"f0_min": args.f0_min,"f0_max": args.f0_max,"stereo_mode": args.stereo_mode}, pipeline_mode="alt" if args.alt_pipeline == True else "orig", embedder_name=args.embedder_name, stack="transformers" if args.use_transformers else "fairseq")
         elif os.path.exists(args.input) and os.path.isdir(args.input):
             list_valid_files = []
             for file in os.listdir(args.input):
@@ -1479,4 +1810,4 @@ if __name__ == "__main__":
             if list_valid_files:
                 for i, vocals_file in enumerate(list_valid_files, start=1):
                     print(f"Файл {i} из {len(list_valid_files)}: {vocals_file}")
-                    vbach_inference(input_file=vocals_file, model_name=args.model_name, output_dir=args.output_dir, output_name=args.output_name, output_bitrate=args.output_bitrate, output_format=args.output_format, pitch=args.pitch, method_pitch=args.method_pitch, format_name=True if len(list_valid_files) > 1 else args.format_name, add_params={ "index_rate": args.index_rate,"filter_radius": args.filter_radius,"protect": args.protect,"rms": args.rms,"mangio_crepe_hop_length": args.hop_length,"f0_min": args.f0_min,"f0_max": args.f0_max,"stereo_mode": args.stereo_mode}, pipeline_mode="alt" if args.alt_pipeline == True else "orig")
+                    vbach_inference(input_file=vocals_file, model_name=args.model_name, output_dir=args.output_dir, output_name=args.output_name, output_bitrate=args.output_bitrate, output_format=args.output_format, pitch=args.pitch, method_pitch=args.method_pitch, format_name=True if len(list_valid_files) > 1 else args.format_name, add_params={ "index_rate": args.index_rate,"filter_radius": args.filter_radius,"protect": args.protect,"rms": args.rms,"mangio_crepe_hop_length": args.hop_length,"f0_min": args.f0_min,"f0_max": args.f0_max,"stereo_mode": args.stereo_mode}, pipeline_mode="alt" if args.alt_pipeline == True else "orig", embedder_name=args.embedder_name, stack="transformers" if args.use_transformers else "fairseq")

mvsepless/vbach_lib/algorithm/attentions.py

@@ -0,0 +1,243 @@
+import math
+import torch
+from .commons import convert_pad_shape
+
+
+class MultiHeadAttention(torch.nn.Module):
+    """
+    Multi-head attention module with optional relative positional encoding and proximal bias.
+
+    Args:
+        channels (int): Number of input channels.
+        out_channels (int): Number of output channels.
+        n_heads (int): Number of attention heads.
+        p_dropout (float, optional): Dropout probability. Defaults to 0.0.
+        window_size (int, optional): Window size for relative positional encoding. Defaults to None.
+        heads_share (bool, optional): Whether to share relative positional embeddings across heads. Defaults to True.
+        block_length (int, optional): Block length for local attention. Defaults to None.
+        proximal_bias (bool, optional): Whether to use proximal bias in self-attention. Defaults to False.
+        proximal_init (bool, optional): Whether to initialize the key projection weights the same as query projection weights. Defaults to False.
+    """
+
+    def __init__(
+        self,
+        channels: int,
+        out_channels: int,
+        n_heads: int,
+        p_dropout: float = 0.0,
+        window_size: int = None,
+        heads_share: bool = True,
+        block_length: int = None,
+        proximal_bias: bool = False,
+        proximal_init: bool = False,
+    ):
+        super().__init__()
+        assert (
+            channels % n_heads == 0
+        ), "Channels must be divisible by the number of heads."
+
+        self.channels = channels
+        self.out_channels = out_channels
+        self.n_heads = n_heads
+        self.k_channels = channels // n_heads
+        self.window_size = window_size
+        self.block_length = block_length
+        self.proximal_bias = proximal_bias
+
+        # Define projections
+        self.conv_q = torch.nn.Conv1d(channels, channels, 1)
+        self.conv_k = torch.nn.Conv1d(channels, channels, 1)
+        self.conv_v = torch.nn.Conv1d(channels, channels, 1)
+        self.conv_o = torch.nn.Conv1d(channels, out_channels, 1)
+
+        self.drop = torch.nn.Dropout(p_dropout)
+
+        # Relative positional encodings
+        if window_size:
+            n_heads_rel = 1 if heads_share else n_heads
+            rel_stddev = self.k_channels**-0.5
+            self.emb_rel_k = torch.nn.Parameter(
+                torch.randn(n_heads_rel, 2 * window_size + 1, self.k_channels)
+                * rel_stddev
+            )
+            self.emb_rel_v = torch.nn.Parameter(
+                torch.randn(n_heads_rel, 2 * window_size + 1, self.k_channels)
+                * rel_stddev
+            )
+
+        # Initialize weights
+        torch.nn.init.xavier_uniform_(self.conv_q.weight)
+        torch.nn.init.xavier_uniform_(self.conv_k.weight)
+        torch.nn.init.xavier_uniform_(self.conv_v.weight)
+        torch.nn.init.xavier_uniform_(self.conv_o.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):
+        # Compute query, key, value projections
+        q, k, v = self.conv_q(x), self.conv_k(c), self.conv_v(c)
+
+        # Compute attention
+        x, self.attn = self.attention(q, k, v, mask=attn_mask)
+
+        # Final output projection
+        return self.conv_o(x)
+
+    def attention(self, query, key, value, mask=None):
+        # Reshape and compute scaled dot-product attention
+        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:
+            assert t_s == t_t, "Relative attention only supports self-attention."
+            scores += self._compute_relative_scores(query, t_s)
+
+        if self.proximal_bias:
+            assert t_s == t_t, "Proximal bias only supports self-attention."
+            scores += self._attention_bias_proximal(t_s).to(scores.device, scores.dtype)
+
+        if mask is not None:
+            scores = scores.masked_fill(mask == 0, -1e4)
+            if self.block_length:
+                block_mask = (
+                    torch.ones_like(scores)
+                    .triu(-self.block_length)
+                    .tril(self.block_length)
+                )
+                scores = scores.masked_fill(block_mask == 0, -1e4)
+
+        # Apply softmax and dropout
+        p_attn = self.drop(torch.nn.functional.softmax(scores, dim=-1))
+
+        # Compute attention output
+        output = torch.matmul(p_attn, value)
+
+        if self.window_size:
+            output += self._apply_relative_values(p_attn, t_s)
+
+        return output.transpose(2, 3).contiguous().view(b, d, t_t), p_attn
+
+    def _compute_relative_scores(self, query, length):
+        rel_emb = self._get_relative_embeddings(self.emb_rel_k, length)
+        rel_logits = self._matmul_with_relative_keys(
+            query / math.sqrt(self.k_channels), rel_emb
+        )
+        return self._relative_position_to_absolute_position(rel_logits)
+
+    def _apply_relative_values(self, p_attn, length):
+        rel_weights = self._absolute_position_to_relative_position(p_attn)
+        rel_emb = self._get_relative_embeddings(self.emb_rel_v, length)
+        return self._matmul_with_relative_values(rel_weights, rel_emb)
+
+    # Helper methods
+    def _matmul_with_relative_values(self, x, y):
+        return torch.matmul(x, y.unsqueeze(0))
+
+    def _matmul_with_relative_keys(self, x, y):
+        return torch.matmul(x, y.unsqueeze(0).transpose(-2, -1))
+
+    def _get_relative_embeddings(self, embeddings, length):
+        pad_length = max(length - (self.window_size + 1), 0)
+        start = max((self.window_size + 1) - length, 0)
+        end = start + 2 * length - 1
+
+        if pad_length > 0:
+            embeddings = torch.nn.functional.pad(
+                embeddings,
+                convert_pad_shape([[0, 0], [pad_length, pad_length], [0, 0]]),
+            )
+        return embeddings[:, start:end]
+
+    def _relative_position_to_absolute_position(self, x):
+        batch, heads, length, _ = x.size()
+        x = torch.nn.functional.pad(
+            x, convert_pad_shape([[0, 0], [0, 0], [0, 0], [0, 1]])
+        )
+        x_flat = x.view(batch, heads, length * 2 * length)
+        x_flat = torch.nn.functional.pad(
+            x_flat, convert_pad_shape([[0, 0], [0, 0], [0, length - 1]])
+        )
+        return x_flat.view(batch, heads, length + 1, 2 * length - 1)[
+            :, :, :length, length - 1 :
+        ]
+
+    def _absolute_position_to_relative_position(self, x):
+        batch, heads, length, _ = x.size()
+        x = torch.nn.functional.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))
+        x_flat = torch.nn.functional.pad(
+            x_flat, convert_pad_shape([[0, 0], [0, 0], [length, 0]])
+        )
+        return x_flat.view(batch, heads, length, 2 * length)[:, :, :, 1:]
+
+    def _attention_bias_proximal(self, length):
+        r = torch.arange(length, dtype=torch.float32)
+        diff = r.unsqueeze(0) - r.unsqueeze(1)
+        return -torch.log1p(torch.abs(diff)).unsqueeze(0).unsqueeze(0)
+
+
+class FFN(torch.nn.Module):
+    """
+    Feed-forward network module.
+
+    Args:
+        in_channels (int): Number of input channels.
+        out_channels (int): Number of output channels.
+        filter_channels (int): Number of filter channels in the convolution layers.
+        kernel_size (int): Kernel size of the convolution layers.
+        p_dropout (float, optional): Dropout probability. Defaults to 0.0.
+        activation (str, optional): Activation function to use. Defaults to None.
+        causal (bool, optional): Whether to use causal padding in the convolution layers. Defaults to False.
+    """
+
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: int,
+        filter_channels: int,
+        kernel_size: int,
+        p_dropout: float = 0.0,
+        activation: str = None,
+        causal: bool = False,
+    ):
+        super().__init__()
+        self.padding_fn = self._causal_padding if causal else self._same_padding
+
+        self.conv_1 = torch.nn.Conv1d(in_channels, filter_channels, kernel_size)
+        self.conv_2 = torch.nn.Conv1d(filter_channels, out_channels, kernel_size)
+        self.drop = torch.nn.Dropout(p_dropout)
+
+        self.activation = activation
+
+    def forward(self, x, x_mask):
+        x = self.conv_1(self.padding_fn(x * x_mask))
+        x = self._apply_activation(x)
+        x = self.drop(x)
+        x = self.conv_2(self.padding_fn(x * x_mask))
+        return x * x_mask
+
+    def _apply_activation(self, x):
+        if self.activation == "gelu":
+            return x * torch.sigmoid(1.702 * x)
+        return torch.relu(x)
+
+    def _causal_padding(self, x):
+        pad_l, pad_r = self.conv_1.kernel_size[0] - 1, 0
+        return torch.nn.functional.pad(
+            x, convert_pad_shape([[0, 0], [0, 0], [pad_l, pad_r]])
+        )
+
+    def _same_padding(self, x):
+        pad = (self.conv_1.kernel_size[0] - 1) // 2
+        return torch.nn.functional.pad(
+            x, convert_pad_shape([[0, 0], [0, 0], [pad, pad]])
+        )

mvsepless/vbach_lib/algorithm/commons.py

@@ -0,0 +1,138 @@
+import torch
+from typing import Optional
+
+
+def init_weights(m, mean=0.0, std=0.01):
+    """
+    Initialize the weights of a module.
+
+    Args:
+        m: The module to initialize.
+        mean: The mean of the normal distribution.
+        std: The standard deviation of the normal distribution.
+    """
+    classname = m.__class__.__name__
+    if classname.find("Conv") != -1:
+        m.weight.data.normal_(mean, std)
+
+
+def get_padding(kernel_size, dilation=1):
+    """
+    Calculate the padding needed for a convolution.
+
+    Args:
+        kernel_size: The size of the kernel.
+        dilation: The dilation of the convolution.
+    """
+    return int((kernel_size * dilation - dilation) / 2)
+
+
+def convert_pad_shape(pad_shape):
+    """
+    Convert the pad shape to a list of integers.
+
+    Args:
+        pad_shape: The pad shape..
+    """
+    l = pad_shape[::-1]
+    pad_shape = [item for sublist in l for item in sublist]
+    return pad_shape
+
+
+def slice_segments(
+    x: torch.Tensor, ids_str: torch.Tensor, segment_size: int = 4, dim: int = 2
+):
+    """
+    Slice segments from a tensor, handling tensors with different numbers of dimensions.
+
+    Args:
+        x (torch.Tensor): The tensor to slice.
+        ids_str (torch.Tensor): The starting indices of the segments.
+        segment_size (int, optional): The size of each segment. Defaults to 4.
+        dim (int, optional): The dimension to slice across (2D or 3D tensors). Defaults to 2.
+    """
+    if dim == 2:
+        ret = torch.zeros_like(x[:, :segment_size])
+    elif dim == 3:
+        ret = torch.zeros_like(x[:, :, :segment_size])
+
+    for i in range(x.size(0)):
+        idx_str = ids_str[i].item()
+        idx_end = idx_str + segment_size
+        if dim == 2:
+            ret[i] = x[i, idx_str:idx_end]
+        else:
+            ret[i] = x[i, :, idx_str:idx_end]
+
+    return ret
+
+
+def rand_slice_segments(x, x_lengths=None, segment_size=4):
+    """
+    Randomly slice segments from a tensor.
+
+    Args:
+        x: The tensor to slice.
+        x_lengths: The lengths of the sequences.
+        segment_size: The size of each segment.
+    """
+    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], device=x.device) * ids_str_max).to(dtype=torch.long)
+    ret = slice_segments(x, ids_str, segment_size, dim=3)
+    return ret, ids_str
+
+
+@torch.jit.script
+def fused_add_tanh_sigmoid_multiply(input_a, input_b, n_channels):
+    """
+    Fused add tanh sigmoid multiply operation.
+
+    Args:
+        input_a: The first input tensor.
+        input_b: The second input tensor.
+        n_channels: The number of 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 sequence_mask(length: torch.Tensor, max_length: Optional[int] = None):
+    """
+    Generate a sequence mask.
+
+    Args:
+        length: The lengths of the sequences.
+        max_length: The maximum length of the sequences.
+    """
+    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 grad_norm(parameters, norm_type: float = 2.0):
+    """
+    Calculates norm of parameter gradients
+
+    Args:
+        parameters: The list of parameters to clip.
+        norm_type: The type of norm to use for clipping.
+    """
+    if isinstance(parameters, torch.Tensor):
+        parameters = [parameters]
+
+    parameters = [p for p in parameters if p.grad is not None]
+
+    if not parameters:
+        return 0.0
+
+    return torch.linalg.vector_norm(
+        torch.stack([p.grad.norm(norm_type) for p in parameters]), ord=norm_type
+    ).item()

mvsepless/vbach_lib/algorithm/discriminators.py

@@ -0,0 +1,262 @@
+import torch
+import torch.nn.functional as F
+from torch.utils.checkpoint import checkpoint
+from torch.nn.utils.parametrizations import spectral_norm, weight_norm
+
+from .commons import get_padding
+from .residuals import LRELU_SLOPE
+
+
+class MultiPeriodDiscriminator(torch.nn.Module):
+    """
+    Multi-period discriminator.
+
+    This class implements a multi-period discriminator, which is used to
+    discriminate between real and fake audio signals. The discriminator
+    is composed of a series of convolutional layers that are applied to
+    the input signal at different periods.
+
+    Args:
+        use_spectral_norm (bool): Whether to use spectral normalization.
+            Defaults to False.
+    """
+
+    def __init__(
+        self,
+        use_spectral_norm: bool = False,
+        checkpointing: bool = False,
+        version: str = "v2",
+    ):
+        super().__init__()
+
+        if version == "v1":
+            periods = [2, 3, 5, 7, 11, 17]
+            resolutions = []
+        elif version == "v2":
+            periods = [2, 3, 5, 7, 11, 17, 23, 37]
+            resolutions = []
+        elif version == "v3":
+            periods = [2, 3, 5, 7, 11]
+            resolutions = [[1024, 120, 600], [2048, 240, 1200], [512, 50, 240]]
+
+        self.checkpointing = checkpointing
+        self.discriminators = torch.nn.ModuleList(
+            [DiscriminatorS(use_spectral_norm=use_spectral_norm)]
+            + [DiscriminatorP(p, use_spectral_norm=use_spectral_norm) for p in periods]
+            + [
+                DiscriminatorR(r, use_spectral_norm=use_spectral_norm)
+                for r in resolutions
+            ]
+        )
+
+    def forward(self, y, y_hat):
+        y_d_rs, y_d_gs, fmap_rs, fmap_gs = [], [], [], []
+        for d in self.discriminators:
+            if self.training and self.checkpointing:
+                y_d_r, fmap_r = checkpoint(d, y, use_reentrant=False)
+                y_d_g, fmap_g = checkpoint(d, y_hat, use_reentrant=False)
+            else:
+                y_d_r, fmap_r = d(y)
+                y_d_g, fmap_g = d(y_hat)
+            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):
+    """
+    Discriminator for the short-term component.
+
+    This class implements a discriminator for the short-term component
+    of the audio signal. The discriminator is composed of a series of
+    convolutional layers that are applied to the input signal.
+    """
+
+    def __init__(self, use_spectral_norm: bool = False):
+        super().__init__()
+
+        norm_f = spectral_norm if use_spectral_norm else weight_norm
+        self.convs = torch.nn.ModuleList(
+            [
+                norm_f(torch.nn.Conv1d(1, 16, 15, 1, padding=7)),
+                norm_f(torch.nn.Conv1d(16, 64, 41, 4, groups=4, padding=20)),
+                norm_f(torch.nn.Conv1d(64, 256, 41, 4, groups=16, padding=20)),
+                norm_f(torch.nn.Conv1d(256, 1024, 41, 4, groups=64, padding=20)),
+                norm_f(torch.nn.Conv1d(1024, 1024, 41, 4, groups=256, padding=20)),
+                norm_f(torch.nn.Conv1d(1024, 1024, 5, 1, padding=2)),
+            ]
+        )
+        self.conv_post = norm_f(torch.nn.Conv1d(1024, 1, 3, 1, padding=1))
+        self.lrelu = torch.nn.LeakyReLU(LRELU_SLOPE)
+
+    def forward(self, x):
+        fmap = []
+        for conv in self.convs:
+            x = self.lrelu(conv(x))
+            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):
+    """
+    Discriminator for the long-term component.
+
+    This class implements a discriminator for the long-term component
+    of the audio signal. The discriminator is composed of a series of
+    convolutional layers that are applied to the input signal at a given
+    period.
+
+    Args:
+        period (int): Period of the discriminator.
+        kernel_size (int): Kernel size of the convolutional layers. Defaults to 5.
+        stride (int): Stride of the convolutional layers. Defaults to 3.
+        use_spectral_norm (bool): Whether to use spectral normalization. Defaults to False.
+    """
+
+    def __init__(
+        self,
+        period: int,
+        kernel_size: int = 5,
+        stride: int = 3,
+        use_spectral_norm: bool = False,
+    ):
+        super().__init__()
+        self.period = period
+        norm_f = spectral_norm if use_spectral_norm else weight_norm
+
+        in_channels = [1, 32, 128, 512, 1024]
+        out_channels = [32, 128, 512, 1024, 1024]
+        strides = [3, 3, 3, 3, 1]
+
+        self.convs = torch.nn.ModuleList(
+            [
+                norm_f(
+                    torch.nn.Conv2d(
+                        in_ch,
+                        out_ch,
+                        (kernel_size, 1),
+                        (s, 1),
+                        padding=(get_padding(kernel_size, 1), 0),
+                    )
+                )
+                for in_ch, out_ch, s in zip(in_channels, out_channels, strides)
+            ]
+        )
+
+        self.conv_post = norm_f(torch.nn.Conv2d(1024, 1, (3, 1), 1, padding=(1, 0)))
+        self.lrelu = torch.nn.LeakyReLU(LRELU_SLOPE)
+
+    def forward(self, x):
+        fmap = []
+        b, c, t = x.shape
+        if t % self.period != 0:
+            n_pad = self.period - (t % self.period)
+            x = torch.nn.functional.pad(x, (0, n_pad), "reflect")
+        x = x.view(b, c, -1, self.period)
+
+        for conv in self.convs:
+            x = self.lrelu(conv(x))
+            fmap.append(x)
+        x = self.conv_post(x)
+        fmap.append(x)
+        x = torch.flatten(x, 1, -1)
+        return x, fmap
+
+
+class DiscriminatorR(torch.nn.Module):
+    def __init__(self, resolution, use_spectral_norm=False):
+        super().__init__()
+
+        self.resolution = resolution
+        self.lrelu_slope = 0.1
+        norm_f = spectral_norm if use_spectral_norm else weight_norm
+
+        self.convs = torch.nn.ModuleList(
+            [
+                norm_f(
+                    torch.nn.Conv2d(
+                        1,
+                        32,
+                        (3, 9),
+                        padding=(1, 4),
+                    )
+                ),
+                norm_f(
+                    torch.nn.Conv2d(
+                        32,
+                        32,
+                        (3, 9),
+                        stride=(1, 2),
+                        padding=(1, 4),
+                    )
+                ),
+                norm_f(
+                    torch.nn.Conv2d(
+                        32,
+                        32,
+                        (3, 9),
+                        stride=(1, 2),
+                        padding=(1, 4),
+                    )
+                ),
+                norm_f(
+                    torch.nn.Conv2d(
+                        32,
+                        32,
+                        (3, 9),
+                        stride=(1, 2),
+                        padding=(1, 4),
+                    )
+                ),
+                norm_f(
+                    torch.nn.Conv2d(
+                        32,
+                        32,
+                        (3, 3),
+                        padding=(1, 1),
+                    )
+                ),
+            ]
+        )
+        self.conv_post = norm_f(torch.nn.Conv2d(32, 1, (3, 3), padding=(1, 1)))
+
+    def forward(self, x):
+        fmap = []
+
+        x = self.spectrogram(x).unsqueeze(1)
+
+        for layer in self.convs:
+            x = F.leaky_relu(layer(x), self.lrelu_slope)
+            fmap.append(x)
+        x = self.conv_post(x)
+        fmap.append(x)
+
+        return torch.flatten(x, 1, -1), fmap
+
+    def spectrogram(self, x):
+        n_fft, hop_length, win_length = self.resolution
+        pad = int((n_fft - hop_length) / 2)
+        x = F.pad(
+            x,
+            (pad, pad),
+            mode="reflect",
+        ).squeeze(1)
+        x = torch.stft(
+            x,
+            n_fft=n_fft,
+            hop_length=hop_length,
+            win_length=win_length,
+            window=torch.ones(win_length, device=x.device),
+            center=False,
+            return_complex=True,
+        )
+
+        mag = torch.norm(torch.view_as_real(x), p=2, dim=-1)  # [B, F, TT]
+
+        return mag

mvsepless/vbach_lib/algorithm/encoders.py

@@ -0,0 +1,209 @@
+import math
+import torch
+from typing import Optional
+
+from .commons import sequence_mask
+from .modules import WaveNet
+from .normalization import LayerNorm
+from .attentions import FFN, MultiHeadAttention
+
+
+class Encoder(torch.nn.Module):
+    """
+    Encoder module for the Transformer model.
+
+    Args:
+        hidden_channels (int): Number of hidden channels in the encoder.
+        filter_channels (int): Number of filter channels in the feed-forward network.
+        n_heads (int): Number of attention heads.
+        n_layers (int): Number of encoder layers.
+        kernel_size (int, optional): Kernel size of the convolution layers in the feed-forward network. Defaults to 1.
+        p_dropout (float, optional): Dropout probability. Defaults to 0.0.
+        window_size (int, optional): Window size for relative positional encoding. Defaults to 10.
+    """
+
+    def __init__(
+        self,
+        hidden_channels: int,
+        filter_channels: int,
+        n_heads: int,
+        n_layers: int,
+        kernel_size: int = 1,
+        p_dropout: float = 0.0,
+        window_size: int = 10,
+    ):
+        super().__init__()
+
+        self.hidden_channels = hidden_channels
+        self.n_layers = n_layers
+        self.drop = torch.nn.Dropout(p_dropout)
+
+        self.attn_layers = torch.nn.ModuleList(
+            [
+                MultiHeadAttention(
+                    hidden_channels,
+                    hidden_channels,
+                    n_heads,
+                    p_dropout=p_dropout,
+                    window_size=window_size,
+                )
+                for _ in range(n_layers)
+            ]
+        )
+        self.norm_layers_1 = torch.nn.ModuleList(
+            [LayerNorm(hidden_channels) for _ in range(n_layers)]
+        )
+        self.ffn_layers = torch.nn.ModuleList(
+            [
+                FFN(
+                    hidden_channels,
+                    hidden_channels,
+                    filter_channels,
+                    kernel_size,
+                    p_dropout=p_dropout,
+                )
+                for _ in range(n_layers)
+            ]
+        )
+        self.norm_layers_2 = torch.nn.ModuleList(
+            [LayerNorm(hidden_channels) for _ in range(n_layers)]
+        )
+
+    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)
+
+        return x * x_mask
+
+
+class TextEncoder(torch.nn.Module):
+    """
+    Text Encoder with configurable embedding dimension.
+
+    Args:
+        out_channels (int): Output channels of the encoder.
+        hidden_channels (int): Hidden channels of the encoder.
+        filter_channels (int): Filter channels of the encoder.
+        n_heads (int): Number of attention heads.
+        n_layers (int): Number of encoder layers.
+        kernel_size (int): Kernel size of the convolutional layers.
+        p_dropout (float): Dropout probability.
+        embedding_dim (int): Embedding dimension for phone embeddings (v1 = 256, v2 = 768).
+        f0 (bool, optional): Whether to use F0 embedding. Defaults to True.
+    """
+
+    def __init__(
+        self,
+        out_channels: int,
+        hidden_channels: int,
+        filter_channels: int,
+        n_heads: int,
+        n_layers: int,
+        kernel_size: int,
+        p_dropout: float,
+        embedding_dim: int,
+        f0: bool = True,
+    ):
+        super().__init__()
+        self.hidden_channels = hidden_channels
+        self.out_channels = out_channels
+        self.emb_phone = torch.nn.Linear(embedding_dim, hidden_channels)
+        self.lrelu = torch.nn.LeakyReLU(0.1, inplace=True)
+        self.emb_pitch = torch.nn.Embedding(256, hidden_channels) if f0 else None
+
+        self.encoder = Encoder(
+            hidden_channels, filter_channels, n_heads, n_layers, kernel_size, p_dropout
+        )
+        self.proj = torch.nn.Conv1d(hidden_channels, out_channels * 2, 1)
+
+    def forward(
+        self, phone: torch.Tensor, pitch: Optional[torch.Tensor], lengths: torch.Tensor
+    ):
+        x = self.emb_phone(phone)
+        if pitch is not None and self.emb_pitch:
+            x += self.emb_pitch(pitch)
+
+        x *= math.sqrt(self.hidden_channels)
+        x = self.lrelu(x)
+        x = x.transpose(1, -1)  # [B, H, T]
+
+        x_mask = sequence_mask(lengths, x.size(2)).unsqueeze(1).to(x.dtype)
+        x = self.encoder(x, x_mask)
+        stats = self.proj(x) * x_mask
+
+        m, logs = torch.split(stats, self.out_channels, dim=1)
+        return m, logs, x_mask
+
+
+class PosteriorEncoder(torch.nn.Module):
+    """
+    Posterior Encoder for inferring latent representation.
+
+    Args:
+        in_channels (int): Number of channels in the input.
+        out_channels (int): Number of channels in the output.
+        hidden_channels (int): Number of hidden channels in the encoder.
+        kernel_size (int): Kernel size of the convolutional layers.
+        dilation_rate (int): Dilation rate of the convolutional layers.
+        n_layers (int): Number of layers in the encoder.
+        gin_channels (int, optional): Number of channels for the global conditioning input. Defaults to 0.
+    """
+
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: int,
+        hidden_channels: int,
+        kernel_size: int,
+        dilation_rate: int,
+        n_layers: int,
+        gin_channels: int = 0,
+    ):
+        super().__init__()
+        self.out_channels = out_channels
+        self.pre = torch.nn.Conv1d(in_channels, hidden_channels, 1)
+        self.enc = WaveNet(
+            hidden_channels,
+            kernel_size,
+            dilation_rate,
+            n_layers,
+            gin_channels=gin_channels,
+        )
+        self.proj = torch.nn.Conv1d(hidden_channels, out_channels * 2, 1)
+
+    def forward(
+        self, x: torch.Tensor, x_lengths: torch.Tensor, g: Optional[torch.Tensor] = None
+    ):
+        x_mask = sequence_mask(x_lengths, x.size(2)).unsqueeze(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)
+        z *= x_mask
+
+        return z, m, logs, x_mask
+
+    def remove_weight_norm(self):
+        self.enc.remove_weight_norm()
+
+    def __prepare_scriptable__(self):
+        for hook in self.enc._forward_pre_hooks.values():
+            if (
+                hook.__module__ == "torch.nn.utils.parametrizations.weight_norm"
+                and hook.__class__.__name__ == "WeightNorm"
+            ):
+                torch.nn.utils.remove_weight_norm(self.enc)
+        return self

mvsepless/vbach_lib/algorithm/generators/hifigan.py

@@ -0,0 +1,228 @@
+import torch
+import numpy as np
+from torch.nn.utils import remove_weight_norm
+from torch.nn.utils.parametrizations import weight_norm
+from typing import Optional
+
+from ..residuals import LRELU_SLOPE, ResBlock
+from ..commons import init_weights
+
+
+class HiFiGANGenerator(torch.nn.Module):
+    """
+    HiFi-GAN Generator module for audio synthesis.
+
+    This module implements the generator part of the HiFi-GAN architecture,
+    which uses transposed convolutions for upsampling and residual blocks for
+    refining the audio output. It can also incorporate global conditioning.
+
+    Args:
+        initial_channel (int): Number of input channels to the initial convolutional layer.
+        resblock_kernel_sizes (list): List of kernel sizes for the residual blocks.
+        resblock_dilation_sizes (list): List of lists of dilation rates for the residual blocks, corresponding to each kernel size.
+        upsample_rates (list): List of upsampling factors for each upsampling layer.
+        upsample_initial_channel (int): Number of output channels from the initial convolutional layer, which is also the input to the first upsampling layer.
+        upsample_kernel_sizes (list): List of kernel sizes for the transposed convolutional layers used for upsampling.
+        gin_channels (int, optional): Number of input channels for the global conditioning. If 0, no global conditioning is used. Defaults to 0.
+    """
+
+    def __init__(
+        self,
+        initial_channel: int,
+        resblock_kernel_sizes: list,
+        resblock_dilation_sizes: list,
+        upsample_rates: list,
+        upsample_initial_channel: int,
+        upsample_kernel_sizes: list,
+        gin_channels: int = 0,
+    ):
+        super(HiFiGANGenerator, self).__init__()
+        self.num_kernels = len(resblock_kernel_sizes)
+        self.num_upsamples = len(upsample_rates)
+        self.conv_pre = torch.nn.Conv1d(
+            initial_channel, upsample_initial_channel, 7, 1, padding=3
+        )
+
+        self.ups = torch.nn.ModuleList()
+        self.resblocks = torch.nn.ModuleList()
+
+        for i, (u, k) in enumerate(zip(upsample_rates, upsample_kernel_sizes)):
+            self.ups.append(
+                weight_norm(
+                    torch.nn.ConvTranspose1d(
+                        upsample_initial_channel // (2**i),
+                        upsample_initial_channel // (2 ** (i + 1)),
+                        k,
+                        u,
+                        padding=(k - u) // 2,
+                    )
+                )
+            )
+            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 = torch.nn.Conv1d(ch, 1, 7, 1, padding=3, bias=False)
+        self.ups.apply(init_weights)
+
+        if gin_channels != 0:
+            self.cond = torch.nn.Conv1d(gin_channels, upsample_initial_channel, 1)
+
+    def forward(self, x: torch.Tensor, g: Optional[torch.Tensor] = None):
+        # new tensor
+        x = self.conv_pre(x)
+
+        if g is not None:
+            x = x + self.cond(g)
+
+        for i in range(self.num_upsamples):
+            x = torch.nn.functional.leaky_relu(x, 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
+        # in-place call
+        x = torch.nn.functional.leaky_relu(x)
+        x = self.conv_post(x)
+        # in-place call
+        x = torch.tanh(x)
+
+        return x
+
+    def __prepare_scriptable__(self):
+        for l in self.ups_and_resblocks:
+            for hook in l._forward_pre_hooks.values():
+                if (
+                    hook.__module__ == "torch.nn.utils.parametrizations.weight_norm"
+                    and hook.__class__.__name__ == "WeightNorm"
+                ):
+                    torch.nn.utils.remove_weight_norm(l)
+        return self
+
+    def remove_weight_norm(self):
+        for l in self.ups:
+            remove_weight_norm(l)
+        for l in self.resblocks:
+            l.remove_weight_norm()
+
+
+class SineGenerator(torch.nn.Module):
+    """
+    Sine wave generator with optional harmonic overtones and noise.
+
+    This module generates sine waves for a fundamental frequency and its harmonics.
+    It can also add Gaussian noise and apply a voiced/unvoiced mask.
+
+    Args:
+        sampling_rate (int): The sampling rate of the audio in Hz.
+        num_harmonics (int, optional): The number of harmonic overtones to generate. Defaults to 0.
+        sine_amplitude (float, optional): The amplitude of the sine wave components. Defaults to 0.1.
+        noise_stddev (float, optional): The standard deviation of the additive Gaussian noise. Defaults to 0.003.
+        voiced_threshold (float, optional): The threshold for the fundamental frequency (F0) to determine if a frame is voiced. Defaults to 0.0.
+    """
+
+    def __init__(
+        self,
+        sampling_rate: int,
+        num_harmonics: int = 0,
+        sine_amplitude: float = 0.1,
+        noise_stddev: float = 0.003,
+        voiced_threshold: float = 0.0,
+    ):
+        super(SineGenerator, self).__init__()
+        self.sampling_rate = sampling_rate
+        self.num_harmonics = num_harmonics
+        self.sine_amplitude = sine_amplitude
+        self.noise_stddev = noise_stddev
+        self.voiced_threshold = voiced_threshold
+        self.waveform_dim = self.num_harmonics + 1  # fundamental + harmonics
+
+    def _compute_voiced_unvoiced(self, f0: torch.Tensor):
+        """
+        Generates a binary mask indicating voiced/unvoiced frames based on the fundamental frequency.
+
+        Args:
+            f0 (torch.Tensor): Fundamental frequency tensor of shape (batch_size, length).
+        """
+        uv_mask = (f0 > self.voiced_threshold).float()
+        return uv_mask
+
+    def _generate_sine_wave(self, f0: torch.Tensor, upsampling_factor: int):
+        """
+        Generates sine waves for the fundamental frequency and its harmonics.
+
+        Args:
+            f0 (torch.Tensor): Fundamental frequency tensor of shape (batch_size, length, 1).
+            upsampling_factor (int): The factor by which to upsample the sine wave.
+        """
+        batch_size, length, _ = f0.shape
+
+        # Create an upsampling grid
+        upsampling_grid = torch.arange(
+            1, upsampling_factor + 1, dtype=f0.dtype, device=f0.device
+        )
+
+        # Calculate phase increments
+        phase_increments = (f0 / self.sampling_rate) * upsampling_grid
+        phase_remainder = torch.fmod(phase_increments[:, :-1, -1:] + 0.5, 1.0) - 0.5
+        cumulative_phase = phase_remainder.cumsum(dim=1).fmod(1.0).to(f0.dtype)
+        phase_increments += torch.nn.functional.pad(
+            cumulative_phase, (0, 0, 1, 0), mode="constant"
+        )
+
+        # Reshape to match the sine wave shape
+        phase_increments = phase_increments.reshape(batch_size, -1, 1)
+
+        # Scale for harmonics
+        harmonic_scale = torch.arange(
+            1, self.waveform_dim + 1, dtype=f0.dtype, device=f0.device
+        ).reshape(1, 1, -1)
+        phase_increments *= harmonic_scale
+
+        # Add random phase offset (except for the fundamental)
+        random_phase = torch.rand(1, 1, self.waveform_dim, device=f0.device)
+        random_phase[..., 0] = 0  # Fundamental frequency has no random offset
+        phase_increments += random_phase
+
+        # Generate sine waves
+        sine_waves = torch.sin(2 * np.pi * phase_increments)
+        return sine_waves
+
+    def forward(self, f0: torch.Tensor, upsampling_factor: int):
+        with torch.no_grad():
+            # Expand `f0` to include waveform dimensions
+            f0 = f0.unsqueeze(-1)
+
+            # Generate sine waves
+            sine_waves = (
+                self._generate_sine_wave(f0, upsampling_factor) * self.sine_amplitude
+            )
+
+            # Compute voiced/unvoiced mask
+            voiced_mask = self._compute_voiced_unvoiced(f0)
+
+            # Upsample voiced/unvoiced mask
+            voiced_mask = torch.nn.functional.interpolate(
+                voiced_mask.transpose(2, 1),
+                scale_factor=float(upsampling_factor),
+                mode="nearest",
+            ).transpose(2, 1)
+
+            # Compute noise amplitude
+            noise_amplitude = voiced_mask * self.noise_stddev + (1 - voiced_mask) * (
+                self.sine_amplitude / 3
+            )
+
+            # Add Gaussian noise
+            noise = noise_amplitude * torch.randn_like(sine_waves)
+
+            # Combine sine waves and noise
+            sine_waveforms = sine_waves * voiced_mask + noise
+
+        return sine_waveforms, voiced_mask, noise

mvsepless/vbach_lib/algorithm/generators/hifigan_mrf.py

@@ -0,0 +1,374 @@
+import math
+from typing import Optional
+
+import numpy as np
+import torch
+from torch.nn.utils import remove_weight_norm
+from torch.nn.utils.parametrizations import weight_norm
+from torch.utils.checkpoint import checkpoint
+
+LRELU_SLOPE = 0.1
+
+
+class MRFLayer(torch.nn.Module):
+    """
+    A single layer of the Multi-Receptive Field (MRF) block.
+
+    This layer consists of two 1D convolutional layers with weight normalization
+    and Leaky ReLU activation in between. The first convolution has a dilation,
+    while the second has a dilation of 1. A skip connection is added from the input
+    to the output.
+
+    Args:
+        channels (int): The number of input and output channels.
+        kernel_size (int): The kernel size of the convolutional layers.
+        dilation (int): The dilation rate for the first convolutional layer.
+    """
+
+    def __init__(self, channels, kernel_size, dilation):
+        super().__init__()
+        self.conv1 = weight_norm(
+            torch.nn.Conv1d(
+                channels,
+                channels,
+                kernel_size,
+                padding=(kernel_size * dilation - dilation) // 2,
+                dilation=dilation,
+            )
+        )
+        self.conv2 = weight_norm(
+            torch.nn.Conv1d(
+                channels, channels, kernel_size, padding=kernel_size // 2, dilation=1
+            )
+        )
+
+    def forward(self, x: torch.Tensor):
+        y = torch.nn.functional.leaky_relu(x, LRELU_SLOPE)
+        y = self.conv1(y)
+        y = torch.nn.functional.leaky_relu(y, LRELU_SLOPE)
+        y = self.conv2(y)
+        return x + y
+
+    def remove_weight_norm(self):
+        remove_weight_norm(self.conv1)
+        remove_weight_norm(self.conv2)
+
+
+class MRFBlock(torch.nn.Module):
+    """
+    A Multi-Receptive Field (MRF) block.
+
+    This block consists of multiple MRFLayers with different dilation rates.
+    It applies each layer sequentially to the input.
+
+    Args:
+        channels (int): The number of input and output channels for the MRFLayers.
+        kernel_size (int): The kernel size for the convolutional layers in the MRFLayers.
+        dilations (list[int]): A list of dilation rates for the MRFLayers.
+    """
+
+    def __init__(self, channels, kernel_size, dilations):
+        super().__init__()
+        self.layers = torch.nn.ModuleList()
+        for dilation in dilations:
+            self.layers.append(MRFLayer(channels, kernel_size, dilation))
+
+    def forward(self, x: torch.Tensor):
+        for layer in self.layers:
+            x = layer(x)
+        return x
+
+    def remove_weight_norm(self):
+        for layer in self.layers:
+            layer.remove_weight_norm()
+
+
+class SineGenerator(torch.nn.Module):
+    """
+    Definition of sine generator
+
+    Generates sine waveforms with optional harmonics and additive noise.
+    Can be used to create harmonic noise source for neural vocoders.
+
+    Args:
+        samp_rate (int): Sampling rate in Hz.
+        harmonic_num (int): Number of harmonic overtones (default 0).
+        sine_amp (float): Amplitude of sine-waveform (default 0.1).
+        noise_std (float): Standard deviation of Gaussian noise (default 0.003).
+        voiced_threshold (float): F0 threshold for voiced/unvoiced classification (default 0).
+    """
+
+    def __init__(
+        self,
+        samp_rate: int,
+        harmonic_num: int = 0,
+        sine_amp: float = 0.1,
+        noise_std: float = 0.003,
+        voiced_threshold: float = 0,
+    ):
+        super(SineGenerator, 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: torch.Tensor):
+        """
+        Generates voiced/unvoiced (UV) signal based on the fundamental frequency (F0).
+
+        Args:
+            f0 (torch.Tensor): Fundamental frequency tensor of shape (batch_size, length, 1).
+        """
+        # generate uv signal
+        uv = torch.ones_like(f0)
+        uv = uv * (f0 > self.voiced_threshold)
+        return uv
+
+    def _f02sine(self, f0_values: torch.Tensor):
+        """
+        Generates sine waveforms based on the fundamental frequency (F0) and its harmonics.
+
+        Args:
+            f0_values (torch.Tensor): Tensor of fundamental frequency and its harmonics,
+                                      shape (batch_size, length, dim), where dim indicates
+                                      the fundamental tone and overtones.
+        """
+        # convert to F0 in rad. The integer part n can be ignored
+        # because 2 * np.pi * n doesn't affect phase
+        rad_values = (f0_values / self.sampling_rate) % 1
+
+        # initial phase noise (no noise for fundamental component)
+        rand_ini = torch.rand(
+            f0_values.shape[0], f0_values.shape[2], device=f0_values.device
+        )
+        rand_ini[:, 0] = 0
+        rad_values[:, 0, :] = rad_values[:, 0, :] + rand_ini
+
+        # instantanouse phase sine[t] = sin(2*pi \sum_i=1 ^{t} rad)
+        tmp_over_one = torch.cumsum(rad_values, 1) % 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
+
+        sines = torch.sin(torch.cumsum(rad_values + cumsum_shift, dim=1) * 2 * np.pi)
+
+        return sines
+
+    def forward(self, f0: torch.Tensor):
+        with torch.no_grad():
+            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)
+
+            sine_waves = self._f02sine(f0_buf) * self.sine_amp
+
+            uv = self._f02uv(f0)
+
+            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):
+    """
+    Generates harmonic and noise source features.
+
+    This module uses the SineGenerator to create harmonic signals based on the
+    fundamental frequency (F0) and merges them into a single excitation signal.
+
+    Args:
+        sample_rate (int): Sampling rate in Hz.
+        harmonic_num (int, optional): Number of harmonics above F0. Defaults to 0.
+        sine_amp (float, optional): Amplitude of sine source signal. Defaults to 0.1.
+        add_noise_std (float, optional): Standard deviation of additive Gaussian noise. Defaults to 0.003.
+        voiced_threshod (float, optional): Threshold to set voiced/unvoiced given F0. Defaults to 0.
+    """
+
+    def __init__(
+        self,
+        sampling_rate: int,
+        harmonic_num: int = 0,
+        sine_amp: float = 0.1,
+        add_noise_std: float = 0.003,
+        voiced_threshold: float = 0,
+    ):
+        super(SourceModuleHnNSF, self).__init__()
+
+        self.sine_amp = sine_amp
+        self.noise_std = add_noise_std
+
+        # to produce sine waveforms
+        self.l_sin_gen = SineGenerator(
+            sampling_rate, harmonic_num, sine_amp, add_noise_std, voiced_threshold
+        )
+
+        # 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: torch.Tensor):
+        sine_wavs, uv, _ = self.l_sin_gen(x)
+        sine_wavs = sine_wavs.to(dtype=self.l_linear.weight.dtype)
+        sine_merge = self.l_tanh(self.l_linear(sine_wavs))
+
+        return sine_merge, None, None
+
+
+class HiFiGANMRFGenerator(torch.nn.Module):
+    """
+    HiFi-GAN generator with Multi-Receptive Field (MRF) blocks.
+
+    This generator takes an input feature sequence and fundamental frequency (F0)
+    as input and generates an audio waveform. It utilizes transposed convolutions
+    for upsampling and MRF blocks for feature refinement. It can also condition
+    on global conditioning features.
+
+    Args:
+        in_channel (int): Number of input channels.
+        upsample_initial_channel (int): Number of channels after the initial convolution.
+        upsample_rates (list[int]): List of upsampling rates for the transposed convolutions.
+        upsample_kernel_sizes (list[int]): List of kernel sizes for the transposed convolutions.
+        resblock_kernel_sizes (list[int]): List of kernel sizes for the convolutional layers in the MRF blocks.
+        resblock_dilations (list[list[int]]): List of lists of dilation rates for the MRF blocks.
+        gin_channels (int): Number of global conditioning input channels (0 if no global conditioning).
+        sample_rate (int): Sampling rate of the audio.
+        harmonic_num (int): Number of harmonics to generate.
+        checkpointing (bool): Whether to use checkpointing to save memory during training (default: False).
+    """
+
+    def __init__(
+        self,
+        in_channel: int,
+        upsample_initial_channel: int,
+        upsample_rates: list[int],
+        upsample_kernel_sizes: list[int],
+        resblock_kernel_sizes: list[int],
+        resblock_dilations: list[list[int]],
+        gin_channels: int,
+        sample_rate: int,
+        harmonic_num: int,
+        checkpointing: bool = False,
+    ):
+        super().__init__()
+        self.num_kernels = len(resblock_kernel_sizes)
+        self.checkpointing = checkpointing
+
+        self.f0_upsample = torch.nn.Upsample(scale_factor=np.prod(upsample_rates))
+        self.m_source = SourceModuleHnNSF(sample_rate, harmonic_num)
+
+        self.conv_pre = weight_norm(
+            torch.nn.Conv1d(
+                in_channel, upsample_initial_channel, kernel_size=7, stride=1, padding=3
+            )
+        )
+        self.upsamples = torch.nn.ModuleList()
+        self.noise_convs = torch.nn.ModuleList()
+
+        stride_f0s = [
+            math.prod(upsample_rates[i + 1 :]) if i + 1 < len(upsample_rates) else 1
+            for i in range(len(upsample_rates))
+        ]
+
+        for i, (u, k) in enumerate(zip(upsample_rates, upsample_kernel_sizes)):
+            # handling odd upsampling rates
+            if u % 2 == 0:
+                # old method
+                padding = (k - u) // 2
+            else:
+                padding = u // 2 + u % 2
+
+            self.upsamples.append(
+                weight_norm(
+                    torch.nn.ConvTranspose1d(
+                        upsample_initial_channel // (2**i),
+                        upsample_initial_channel // (2 ** (i + 1)),
+                        kernel_size=k,
+                        stride=u,
+                        padding=padding,
+                        output_padding=u % 2,
+                    )
+                )
+            )
+            """ handling odd upsampling rates
+            #  s   k   p
+            # 40  80  20
+            # 32  64  16
+            #  4   8   2
+            #  2   3   1
+            # 63 125  31
+            #  9  17   4
+            #  3   5   1
+            #  1   1   0
+            """
+            stride = stride_f0s[i]
+            kernel = 1 if stride == 1 else stride * 2 - stride % 2
+            padding = 0 if stride == 1 else (kernel - stride) // 2
+
+            self.noise_convs.append(
+                torch.nn.Conv1d(
+                    1,
+                    upsample_initial_channel // (2 ** (i + 1)),
+                    kernel_size=kernel,
+                    stride=stride,
+                    padding=padding,
+                )
+            )
+        self.mrfs = torch.nn.ModuleList()
+        for i in range(len(self.upsamples)):
+            channel = upsample_initial_channel // (2 ** (i + 1))
+            self.mrfs.append(
+                torch.nn.ModuleList(
+                    [
+                        MRFBlock(channel, kernel_size=k, dilations=d)
+                        for k, d in zip(resblock_kernel_sizes, resblock_dilations)
+                    ]
+                )
+            )
+        self.conv_post = weight_norm(
+            torch.nn.Conv1d(channel, 1, kernel_size=7, stride=1, padding=3)
+        )
+        if gin_channels != 0:
+            self.cond = torch.nn.Conv1d(gin_channels, upsample_initial_channel, 1)
+
+    def forward(
+        self, x: torch.Tensor, f0: torch.Tensor, g: Optional[torch.Tensor] = None
+    ):
+        f0 = self.f0_upsample(f0[:, None, :]).transpose(-1, -2)
+        har_source, _, _ = self.m_source(f0)
+        har_source = har_source.transpose(-1, -2)
+        x = self.conv_pre(x)
+
+        if g is not None:
+            x = x + self.cond(g)
+
+        for ups, mrf, noise_conv in zip(self.upsamples, self.mrfs, self.noise_convs):
+            x = torch.nn.functional.leaky_relu(x, LRELU_SLOPE)
+
+            if self.training and self.checkpointing:
+                x = checkpoint(ups, x, use_reentrant=False)
+                x = x + noise_conv(har_source)
+                xs = sum([checkpoint(layer, x, use_reentrant=False) for layer in mrf])
+            else:
+                x = ups(x)
+                x = x + noise_conv(har_source)
+                xs = sum([layer(x) for layer in mrf])
+            x = xs / self.num_kernels
+
+        x = torch.nn.functional.leaky_relu(x)
+        x = torch.tanh(self.conv_post(x))
+
+        return x
+
+    def remove_weight_norm(self):
+        remove_weight_norm(self.conv_pre)
+        for up in self.upsamples:
+            remove_weight_norm(up)
+        for mrf in self.mrfs:
+            mrf.remove_weight_norm()
+        remove_weight_norm(self.conv_post)

mvsepless/vbach_lib/algorithm/generators/hifigan_nsf.py

@@ -0,0 +1,235 @@
+import math
+from typing import Optional
+
+import torch
+from torch.nn.utils import remove_weight_norm
+from torch.nn.utils.parametrizations import weight_norm
+from torch.utils.checkpoint import checkpoint
+
+from ..commons import init_weights
+from .hifigan import SineGenerator
+from ..residuals import LRELU_SLOPE, ResBlock
+
+
+class SourceModuleHnNSF(torch.nn.Module):
+    """
+    Source Module for generating harmonic and noise components for audio synthesis.
+
+    This module generates a harmonic source signal using sine waves and adds
+    optional noise. It's often used in neural vocoders as a source of excitation.
+
+    Args:
+        sample_rate (int): Sampling rate of the audio in Hz.
+        harmonic_num (int, optional): Number of harmonic overtones to generate above the fundamental frequency (F0). Defaults to 0.
+        sine_amp (float, optional): Amplitude of the sine wave components. Defaults to 0.1.
+        add_noise_std (float, optional): Standard deviation of the additive white Gaussian noise. Defaults to 0.003.
+        voiced_threshod (float, optional): Threshold for the fundamental frequency (F0) to determine if a frame is voiced. If F0 is below this threshold, it's considered unvoiced. Defaults to 0.
+    """
+
+    def __init__(
+        self,
+        sample_rate: int,
+        harmonic_num: int = 0,
+        sine_amp: float = 0.1,
+        add_noise_std: float = 0.003,
+        voiced_threshod: float = 0,
+    ):
+        super(SourceModuleHnNSF, self).__init__()
+
+        self.sine_amp = sine_amp
+        self.noise_std = add_noise_std
+
+        self.l_sin_gen = SineGenerator(
+            sample_rate, harmonic_num, sine_amp, add_noise_std, voiced_threshod
+        )
+        self.l_linear = torch.nn.Linear(harmonic_num + 1, 1)
+        self.l_tanh = torch.nn.Tanh()
+
+    def forward(self, x: torch.Tensor, upsample_factor: int = 1):
+        sine_wavs, uv, _ = self.l_sin_gen(x, upsample_factor)
+        sine_wavs = sine_wavs.to(dtype=self.l_linear.weight.dtype)
+        sine_merge = self.l_tanh(self.l_linear(sine_wavs))
+        return sine_merge, None, None
+
+
+class HiFiGANNSFGenerator(torch.nn.Module):
+    """
+    Generator module based on the Neural Source Filter (NSF) architecture.
+
+    This generator synthesizes audio by first generating a source excitation signal
+    (harmonic and noise) and then filtering it through a series of upsampling and
+    residual blocks. Global conditioning can be applied to influence the generation.
+
+    Args:
+        initial_channel (int): Number of input channels to the initial convolutional layer.
+        resblock_kernel_sizes (list): List of kernel sizes for the residual blocks.
+        resblock_dilation_sizes (list): List of lists of dilation rates for the residual blocks, corresponding to each kernel size.
+        upsample_rates (list): List of upsampling factors for each upsampling layer.
+        upsample_initial_channel (int): Number of output channels from the initial convolutional layer, which is also the input to the first upsampling layer.
+        upsample_kernel_sizes (list): List of kernel sizes for the transposed convolutional layers used for upsampling.
+        gin_channels (int): Number of input channels for the global conditioning. If 0, no global conditioning is used.
+        sr (int): Sampling rate of the audio.
+        checkpointing (bool, optional): Whether to use gradient checkpointing to save memory during training. Defaults to False.
+    """
+
+    def __init__(
+        self,
+        initial_channel: int,
+        resblock_kernel_sizes: list,
+        resblock_dilation_sizes: list,
+        upsample_rates: list,
+        upsample_initial_channel: int,
+        upsample_kernel_sizes: list,
+        gin_channels: int,
+        sr: int,
+        checkpointing: bool = False,
+    ):
+        super(HiFiGANNSFGenerator, self).__init__()
+
+        self.num_kernels = len(resblock_kernel_sizes)
+        self.num_upsamples = len(upsample_rates)
+        self.checkpointing = checkpointing
+        self.f0_upsamp = torch.nn.Upsample(scale_factor=math.prod(upsample_rates))
+        self.m_source = SourceModuleHnNSF(sample_rate=sr, harmonic_num=0)
+
+        self.conv_pre = torch.nn.Conv1d(
+            initial_channel, upsample_initial_channel, 7, 1, padding=3
+        )
+
+        self.ups = torch.nn.ModuleList()
+        self.noise_convs = torch.nn.ModuleList()
+
+        channels = [
+            upsample_initial_channel // (2 ** (i + 1))
+            for i in range(len(upsample_rates))
+        ]
+        stride_f0s = [
+            math.prod(upsample_rates[i + 1 :]) if i + 1 < len(upsample_rates) else 1
+            for i in range(len(upsample_rates))
+        ]
+
+        for i, (u, k) in enumerate(zip(upsample_rates, upsample_kernel_sizes)):
+            # handling odd upsampling rates
+            if u % 2 == 0:
+                # old method
+                padding = (k - u) // 2
+            else:
+                padding = u // 2 + u % 2
+
+            self.ups.append(
+                weight_norm(
+                    torch.nn.ConvTranspose1d(
+                        upsample_initial_channel // (2**i),
+                        channels[i],
+                        k,
+                        u,
+                        padding=padding,
+                        output_padding=u % 2,
+                    )
+                )
+            )
+            """ handling odd upsampling rates
+            #  s   k   p
+            # 40  80  20
+            # 32  64  16
+            #  4   8   2
+            #  2   3   1
+            # 63 125  31
+            #  9  17   4
+            #  3   5   1
+            #  1   1   0
+            """
+            stride = stride_f0s[i]
+            kernel = 1 if stride == 1 else stride * 2 - stride % 2
+            padding = 0 if stride == 1 else (kernel - stride) // 2
+
+            self.noise_convs.append(
+                torch.nn.Conv1d(
+                    1,
+                    channels[i],
+                    kernel_size=kernel,
+                    stride=stride,
+                    padding=padding,
+                )
+            )
+
+        self.resblocks = torch.nn.ModuleList(
+            [
+                ResBlock(channels[i], k, d)
+                for i in range(len(self.ups))
+                for k, d in zip(resblock_kernel_sizes, resblock_dilation_sizes)
+            ]
+        )
+
+        self.conv_post = torch.nn.Conv1d(channels[-1], 1, 7, 1, padding=3, bias=False)
+        self.ups.apply(init_weights)
+
+        if gin_channels != 0:
+            self.cond = torch.nn.Conv1d(gin_channels, upsample_initial_channel, 1)
+
+        self.upp = math.prod(upsample_rates)
+        self.lrelu_slope = LRELU_SLOPE
+
+    def forward(
+        self, x: torch.Tensor, f0: torch.Tensor, g: Optional[torch.Tensor] = None
+    ):
+        har_source, _, _ = self.m_source(f0, self.upp)
+        har_source = har_source.transpose(1, 2)
+        # new tensor
+        x = self.conv_pre(x)
+
+        if g is not None:
+            x = x + self.cond(g)
+
+        for i, (ups, noise_convs) in enumerate(zip(self.ups, self.noise_convs)):
+            x = torch.nn.functional.leaky_relu(x, self.lrelu_slope)
+            # Apply upsampling layer
+            if self.training and self.checkpointing:
+                x = checkpoint(ups, x, use_reentrant=False)
+                x = x + noise_convs(har_source)
+                xs = sum(
+                    [
+                        checkpoint(resblock, x, use_reentrant=False)
+                        for j, resblock in enumerate(self.resblocks)
+                        if j in range(i * self.num_kernels, (i + 1) * self.num_kernels)
+                    ]
+                )
+            else:
+                x = ups(x)
+                x = x + noise_convs(har_source)
+                xs = sum(
+                    [
+                        resblock(x)
+                        for j, resblock in enumerate(self.resblocks)
+                        if j in range(i * self.num_kernels, (i + 1) * self.num_kernels)
+                    ]
+                )
+            x = xs / self.num_kernels
+
+        x = torch.nn.functional.leaky_relu(x)
+        x = torch.tanh(self.conv_post(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()
+
+    def __prepare_scriptable__(self):
+        for l in self.ups:
+            for hook in l._forward_pre_hooks.values():
+                if (
+                    hook.__module__ == "torch.nn.utils.parametrizations.weight_norm"
+                    and hook.__class__.__name__ == "WeightNorm"
+                ):
+                    remove_weight_norm(l)
+        for l in self.resblocks:
+            for hook in l._forward_pre_hooks.values():
+                if (
+                    hook.__module__ == "torch.nn.utils.parametrizations.weight_norm"
+                    and hook.__class__.__name__ == "WeightNorm"
+                ):
+                    remove_weight_norm(l)
+        return self

mvsepless/vbach_lib/algorithm/generators/refinegan.py

@@ -0,0 +1,451 @@
+import numpy as np
+import torch
+import torchaudio
+from torch import nn
+from torch.nn import functional as F
+from torch.nn.utils.parametrizations import weight_norm
+from torch.nn.utils import remove_weight_norm
+from torch.utils.checkpoint import checkpoint
+
+from ..commons import init_weights, get_padding
+
+
+class ResBlock(nn.Module):
+    """
+    Residual block with multiple dilated convolutions.
+
+    This block applies a sequence of dilated convolutional layers with Leaky ReLU activation.
+    It's designed to capture information at different scales due to the varying dilation rates.
+
+    Args:
+        in_channels (int): Number of input channels.
+        out_channels (int): Number of output channels.
+        kernel_size (int, optional): Kernel size for the convolutional layers. Defaults to 7.
+        dilation (tuple[int], optional): Tuple of dilation rates for the convolutional layers. Defaults to (1, 3, 5).
+        leaky_relu_slope (float, optional): Slope for the Leaky ReLU activation. Defaults to 0.2.
+    """
+
+    def __init__(
+        self,
+        channels: int,
+        kernel_size: int = 7,
+        dilation: tuple[int] = (1, 3, 5),
+        leaky_relu_slope: float = 0.2,
+    ):
+        super().__init__()
+
+        self.leaky_relu_slope = leaky_relu_slope
+
+        self.convs1 = nn.ModuleList(
+            [
+                weight_norm(
+                    nn.Conv1d(
+                        channels,
+                        channels,
+                        kernel_size,
+                        stride=1,
+                        dilation=d,
+                        padding=get_padding(kernel_size, d),
+                    )
+                )
+                for d in dilation
+            ]
+        )
+        self.convs1.apply(init_weights)
+
+        self.convs2 = nn.ModuleList(
+            [
+                weight_norm(
+                    nn.Conv1d(
+                        channels,
+                        channels,
+                        kernel_size,
+                        stride=1,
+                        dilation=1,
+                        padding=get_padding(kernel_size, 1),
+                    )
+                )
+                for d in dilation
+            ]
+        )
+        self.convs2.apply(init_weights)
+
+    def forward(self, x: torch.Tensor):
+        for c1, c2 in zip(self.convs1, self.convs2):
+            xt = F.leaky_relu(x, self.leaky_relu_slope)
+            xt = c1(xt)
+            xt = F.leaky_relu(xt, self.leaky_relu_slope)
+            xt = c2(xt)
+            x = xt + x
+
+        return x
+
+    def remove_weight_norm(self):
+        for c1, c2 in zip(self.convs1, self.convs2):
+            remove_weight_norm(c1)
+            remove_weight_norm(c2)
+
+
+class AdaIN(nn.Module):
+    """
+    Adaptive Instance Normalization layer.
+
+    This layer applies a scaling factor to the input based on a learnable weight.
+
+    Args:
+        channels (int): Number of input channels.
+        leaky_relu_slope (float, optional): Slope for the Leaky ReLU activation applied after scaling. Defaults to 0.2.
+    """
+
+    def __init__(
+        self,
+        *,
+        channels: int,
+        leaky_relu_slope: float = 0.2,
+    ):
+        super().__init__()
+
+        self.weight = nn.Parameter(torch.ones(channels) * 1e-4)
+        # safe to use in-place as it is used on a new x+gaussian tensor
+        self.activation = nn.LeakyReLU(leaky_relu_slope)
+
+    def forward(self, x: torch.Tensor):
+        gaussian = torch.randn_like(x) * self.weight[None, :, None]
+
+        return self.activation(x + gaussian)
+
+
+class ParallelResBlock(nn.Module):
+    """
+    Parallel residual block that applies multiple residual blocks with different kernel sizes in parallel.
+
+    Args:
+        in_channels (int): Number of input channels.
+        out_channels (int): Number of output channels.
+        kernel_sizes (tuple[int], optional): Tuple of kernel sizes for the parallel residual blocks. Defaults to (3, 7, 11).
+        dilation (tuple[int], optional): Tuple of dilation rates for the convolutional layers within the residual blocks. Defaults to (1, 3, 5).
+        leaky_relu_slope (float, optional): Slope for the Leaky ReLU activation. Defaults to 0.2.
+    """
+
+    def __init__(
+        self,
+        *,
+        in_channels: int,
+        out_channels: int,
+        kernel_sizes: tuple[int] = (3, 7, 11),
+        dilation: tuple[int] = (1, 3, 5),
+        leaky_relu_slope: float = 0.2,
+    ):
+        super().__init__()
+
+        self.in_channels = in_channels
+        self.out_channels = out_channels
+
+        self.input_conv = nn.Conv1d(
+            in_channels=in_channels,
+            out_channels=out_channels,
+            kernel_size=7,
+            stride=1,
+            padding=3,
+        )
+
+        self.input_conv.apply(init_weights)
+
+        self.blocks = nn.ModuleList(
+            [
+                nn.Sequential(
+                    AdaIN(channels=out_channels),
+                    ResBlock(
+                        out_channels,
+                        kernel_size=kernel_size,
+                        dilation=dilation,
+                        leaky_relu_slope=leaky_relu_slope,
+                    ),
+                    AdaIN(channels=out_channels),
+                )
+                for kernel_size in kernel_sizes
+            ]
+        )
+
+    def forward(self, x: torch.Tensor):
+        x = self.input_conv(x)
+        return torch.stack([block(x) for block in self.blocks], dim=0).mean(dim=0)
+
+    def remove_weight_norm(self):
+        remove_weight_norm(self.input_conv)
+        for block in self.blocks:
+            block[1].remove_weight_norm()
+
+
+class SineGenerator(nn.Module):
+    """
+    Definition of sine generator
+
+    Generates sine waveforms with optional harmonics and additive noise.
+    Can be used to create harmonic noise source for neural vocoders.
+
+    Args:
+        samp_rate (int): Sampling rate in Hz.
+        harmonic_num (int): Number of harmonic overtones (default 0).
+        sine_amp (float): Amplitude of sine-waveform (default 0.1).
+        noise_std (float): Standard deviation of Gaussian noise (default 0.003).
+        voiced_threshold (float): F0 threshold for voiced/unvoiced classification (default 0).
+    """
+
+    def __init__(
+        self,
+        samp_rate,
+        harmonic_num=0,
+        sine_amp=0.1,
+        noise_std=0.003,
+        voiced_threshold=0,
+    ):
+        super(SineGenerator, 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
+
+        self.merge = nn.Sequential(
+            nn.Linear(self.dim, 1, bias=False),
+            nn.Tanh(),
+        )
+
+    def _f02uv(self, f0):
+        # generate uv signal
+        uv = torch.ones_like(f0)
+        uv = uv * (f0 > self.voiced_threshold)
+        return uv
+
+    def _f02sine(self, f0_values):
+        """f0_values: (batchsize, length, dim)
+        where dim indicates fundamental tone and overtones
+        """
+        # convert to F0 in rad. The integer part n can be ignored
+        # because 2 * np.pi * n doesn't affect phase
+        rad_values = (f0_values / self.sampling_rate) % 1
+
+        # initial phase noise (no noise for fundamental component)
+        rand_ini = torch.rand(
+            f0_values.shape[0], f0_values.shape[2], device=f0_values.device
+        )
+        rand_ini[:, 0] = 0
+        rad_values[:, 0, :] = rad_values[:, 0, :] + rand_ini
+
+        # instantanouse phase sine[t] = sin(2*pi \sum_i=1 ^{t} rad)
+        tmp_over_one = torch.cumsum(rad_values, 1) % 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
+
+        sines = torch.sin(torch.cumsum(rad_values + cumsum_shift, dim=1) * 2 * np.pi)
+
+        return sines
+
+    def forward(self, f0):
+        with torch.no_grad():
+            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)
+
+            sine_waves = self._f02sine(f0_buf) * self.sine_amp
+
+            uv = self._f02uv(f0)
+
+            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
+
+        # merge with grad
+        return self.merge(sine_waves)
+
+
+class RefineGANGenerator(nn.Module):
+    """
+    RefineGAN generator for audio synthesis.
+
+    This generator uses a combination of downsampling, residual blocks, and parallel residual blocks
+    to refine an input mel-spectrogram and fundamental frequency (F0) into an audio waveform.
+    It can also incorporate global conditioning.
+
+    Args:
+        sample_rate (int, optional): Sampling rate of the audio. Defaults to 44100.
+        downsample_rates (tuple[int], optional): Downsampling rates for the downsampling blocks. Defaults to (2, 2, 8, 8).
+        upsample_rates (tuple[int], optional): Upsampling rates for the upsampling blocks. Defaults to (8, 8, 2, 2).
+        leaky_relu_slope (float, optional): Slope for the Leaky ReLU activation. Defaults to 0.2.
+        num_mels (int, optional): Number of mel-frequency bins in the input mel-spectrogram. Defaults to 128.
+        start_channels (int, optional): Number of channels in the initial convolutional layer. Defaults to 16.
+        gin_channels (int, optional): Number of channels for the global conditioning input. Defaults to 256.
+        checkpointing (bool, optional): Whether to use checkpointing for memory efficiency. Defaults to False.
+    """
+
+    def __init__(
+        self,
+        *,
+        sample_rate: int = 44100,
+        downsample_rates: tuple[int] = (2, 2, 8, 8),  # unused
+        upsample_rates: tuple[int] = (8, 8, 2, 2),
+        leaky_relu_slope: float = 0.2,
+        num_mels: int = 128,
+        start_channels: int = 16,  # unused
+        gin_channels: int = 256,
+        checkpointing: bool = False,
+        upsample_initial_channel=512,
+    ):
+        super().__init__()
+        self.upsample_rates = upsample_rates
+        self.leaky_relu_slope = leaky_relu_slope
+        self.checkpointing = checkpointing
+
+        self.upp = np.prod(upsample_rates)
+        self.m_source = SineGenerator(sample_rate)
+
+        # expanded f0 sinegen -> match mel_conv
+        # (8, 1, 17280) -> (8, 16, 17280)
+        self.pre_conv = weight_norm(
+            nn.Conv1d(
+                1,
+                16,
+                7,
+                1,
+                padding=3,
+            )
+        )
+
+        # (8,  16, 17280) = 4th upscale
+        # (8,  32, 8640)  = 3rd upscale
+        # (8,  64, 4320)  = 2nd upscale
+        # (8, 128, 432)   = 1st upscale
+        # (8, 256, 36) merged to mel
+
+        # f0 downsampling and upchanneling
+        channels = start_channels
+        size = self.upp
+        self.downsample_blocks = nn.ModuleList([])
+        self.df0 = []
+        for i, u in enumerate(upsample_rates):
+
+            new_size = int(size / upsample_rates[-i - 1])
+            # T dimension factors for torchaudio.functional.resample
+            self.df0.append([size, new_size])
+            size = new_size
+
+            new_channels = channels * 2
+            self.downsample_blocks.append(
+                weight_norm(nn.Conv1d(channels, new_channels, 7, 1, padding=3))
+            )
+            channels = new_channels
+
+        # mel handling
+        channels = upsample_initial_channel
+
+        self.mel_conv = weight_norm(
+            nn.Conv1d(
+                num_mels,
+                channels // 2,
+                7,
+                1,
+                padding=3,
+            )
+        )
+
+        self.mel_conv.apply(init_weights)
+
+        if gin_channels != 0:
+            self.cond = nn.Conv1d(256, channels // 2, 1)
+
+        self.upsample_blocks = nn.ModuleList([])
+        self.upsample_conv_blocks = nn.ModuleList([])
+
+        for rate in upsample_rates:
+            new_channels = channels // 2
+
+            self.upsample_blocks.append(nn.Upsample(scale_factor=rate, mode="linear"))
+
+            self.upsample_conv_blocks.append(
+                ParallelResBlock(
+                    in_channels=channels + channels // 4,
+                    out_channels=new_channels,
+                    kernel_sizes=(3, 7, 11),
+                    dilation=(1, 3, 5),
+                    leaky_relu_slope=leaky_relu_slope,
+                )
+            )
+
+            channels = new_channels
+
+        self.conv_post = weight_norm(
+            nn.Conv1d(channels, 1, 7, 1, padding=3, bias=False)
+        )
+        self.conv_post.apply(init_weights)
+
+    def forward(self, mel: torch.Tensor, f0: torch.Tensor, g: torch.Tensor = None):
+        f0_size = mel.shape[-1]
+        # change f0 helper to full size
+        f0 = F.interpolate(f0.unsqueeze(1), size=f0_size * self.upp, mode="linear")
+        # get f0 turned into sines harmonics
+        har_source = self.m_source(f0.transpose(1, 2)).transpose(1, 2)
+        # prepare for fusion to mel
+        x = self.pre_conv(har_source)
+        # downsampled/upchanneled versions for each upscale
+        downs = []
+        for block, (old_size, new_size) in zip(self.downsample_blocks, self.df0):
+            x = F.leaky_relu(x, self.leaky_relu_slope)
+            downs.append(x)
+            # attempt to cancel spectral aliasing
+            x = torchaudio.functional.resample(
+                x.contiguous(),
+                orig_freq=int(f0_size * old_size),
+                new_freq=int(f0_size * new_size),
+                lowpass_filter_width=64,
+                rolloff=0.9475937167399596,
+                resampling_method="sinc_interp_kaiser",
+                beta=14.769656459379492,
+            )
+            x = block(x)
+
+        # expanding spectrogram from 192 to 256 channels
+        mel = self.mel_conv(mel)
+        if g is not None:
+            # adding expanded speaker embedding
+            mel = mel + self.cond(g)
+
+        x = torch.cat([mel, x], dim=1)
+
+        for ups, res, down in zip(
+            self.upsample_blocks,
+            self.upsample_conv_blocks,
+            reversed(downs),
+        ):
+            x = F.leaky_relu(x, self.leaky_relu_slope)
+
+            if self.training and self.checkpointing:
+                x = checkpoint(ups, x, use_reentrant=False)
+                x = torch.cat([x, down], dim=1)
+                x = checkpoint(res, x, use_reentrant=False)
+            else:
+                x = ups(x)
+                x = torch.cat([x, down], dim=1)
+                x = res(x)
+
+        x = F.leaky_relu(x, self.leaky_relu_slope)
+        x = self.conv_post(x)
+        x = torch.tanh(x)
+
+        return x
+
+    def remove_weight_norm(self):
+        remove_weight_norm(self.pre_conv)
+        remove_weight_norm(self.mel_conv)
+        remove_weight_norm(self.conv_post)
+
+        for block in self.downsample_blocks:
+            block.remove_weight_norm()
+
+        for block in self.upsample_conv_blocks:
+            block.remove_weight_norm()

mvsepless/vbach_lib/algorithm/modules.py

@@ -0,0 +1,117 @@
+import torch
+from .commons import fused_add_tanh_sigmoid_multiply
+
+
+class WaveNet(torch.nn.Module):
+    """
+    WaveNet residual blocks as used in WaveGlow.
+
+    Args:
+        hidden_channels (int): Number of hidden channels.
+        kernel_size (int): Size of the convolutional kernel.
+        dilation_rate (int): Dilation rate of the convolution.
+        n_layers (int): Number of convolutional layers.
+        gin_channels (int, optional): Number of conditioning channels. Defaults to 0.
+        p_dropout (float, optional): Dropout probability. Defaults to 0.
+    """
+
+    def __init__(
+        self,
+        hidden_channels: int,
+        kernel_size: int,
+        dilation_rate,
+        n_layers: int,
+        gin_channels: int = 0,
+        p_dropout: int = 0,
+    ):
+        super().__init__()
+        assert kernel_size % 2 == 1, "Kernel size must be odd for proper padding."
+
+        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.n_channels_tensor = torch.IntTensor([hidden_channels])  # Static tensor
+
+        self.in_layers = torch.nn.ModuleList()
+        self.res_skip_layers = torch.nn.ModuleList()
+        self.drop = torch.nn.Dropout(p_dropout)
+
+        # Conditional layer for global conditioning
+        if gin_channels:
+            self.cond_layer = torch.nn.utils.parametrizations.weight_norm(
+                torch.nn.Conv1d(gin_channels, 2 * hidden_channels * n_layers, 1),
+                name="weight",
+            )
+
+        # Precompute dilations and paddings
+        dilations = [dilation_rate**i for i in range(n_layers)]
+        paddings = [(kernel_size * d - d) // 2 for d in dilations]
+
+        # Initialize layers
+        for i in range(n_layers):
+            self.in_layers.append(
+                torch.nn.utils.parametrizations.weight_norm(
+                    torch.nn.Conv1d(
+                        hidden_channels,
+                        2 * hidden_channels,
+                        kernel_size,
+                        dilation=dilations[i],
+                        padding=paddings[i],
+                    ),
+                    name="weight",
+                )
+            )
+
+            res_skip_channels = (
+                hidden_channels if i == n_layers - 1 else 2 * hidden_channels
+            )
+            self.res_skip_layers.append(
+                torch.nn.utils.parametrizations.weight_norm(
+                    torch.nn.Conv1d(hidden_channels, res_skip_channels, 1),
+                    name="weight",
+                )
+            )
+
+    def forward(self, x, x_mask, g=None):
+        output = x.clone().zero_()
+
+        # Apply conditional layer if global conditioning is provided
+        g = self.cond_layer(g) if g is not None else None
+
+        for i in range(self.n_layers):
+            x_in = self.in_layers[i](x)
+            g_l = (
+                g[
+                    :,
+                    i * 2 * self.hidden_channels : (i + 1) * 2 * self.hidden_channels,
+                    :,
+                ]
+                if g is not None
+                else 0
+            )
+
+            # Activation with fused Tanh-Sigmoid
+            acts = fused_add_tanh_sigmoid_multiply(x_in, g_l, self.n_channels_tensor)
+            acts = self.drop(acts)
+
+            # Residual and skip connections
+            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:
+            torch.nn.utils.remove_weight_norm(self.cond_layer)
+        for layer in self.in_layers:
+            torch.nn.utils.remove_weight_norm(layer)
+        for layer in self.res_skip_layers:
+            torch.nn.utils.remove_weight_norm(layer)

mvsepless/vbach_lib/algorithm/normalization.py

@@ -0,0 +1,26 @@
+import torch
+
+
+class LayerNorm(torch.nn.Module):
+    """
+    Layer normalization module.
+
+    Args:
+        channels (int): Number of channels.
+        eps (float, optional): Epsilon value for numerical stability. Defaults to 1e-5.
+    """
+
+    def __init__(self, channels: int, eps: float = 1e-5):
+        super().__init__()
+        self.eps = eps
+        self.gamma = torch.nn.Parameter(torch.ones(channels))
+        self.beta = torch.nn.Parameter(torch.zeros(channels))
+
+    def forward(self, x):
+        # Transpose to (batch_size, time_steps, channels) for layer_norm
+        x = x.transpose(1, -1)
+        x = torch.nn.functional.layer_norm(
+            x, (x.size(-1),), self.gamma, self.beta, self.eps
+        )
+        # Transpose back to (batch_size, channels, time_steps)
+        return x.transpose(1, -1)

mvsepless/vbach_lib/algorithm/residuals.py

@@ -0,0 +1,261 @@
+import torch
+from itertools import chain
+from typing import Optional, Tuple
+from torch.nn.utils import remove_weight_norm
+from torch.nn.utils.parametrizations import weight_norm
+
+from .modules import WaveNet
+from .commons import get_padding, init_weights
+
+LRELU_SLOPE = 0.1
+
+
+def create_conv1d_layer(channels, kernel_size, dilation):
+    return weight_norm(
+        torch.nn.Conv1d(
+            channels,
+            channels,
+            kernel_size,
+            1,
+            dilation=dilation,
+            padding=get_padding(kernel_size, dilation),
+        )
+    )
+
+
+def apply_mask(tensor: torch.Tensor, mask: Optional[torch.Tensor]):
+    return tensor * mask if mask else tensor
+
+
+def apply_mask_(tensor: torch.Tensor, mask: Optional[torch.Tensor]):
+    return tensor.mul_(mask) if mask else tensor
+
+
+class ResBlock(torch.nn.Module):
+    """
+    A residual block module that applies a series of 1D convolutional layers with residual connections.
+    """
+
+    def __init__(
+        self, channels: int, kernel_size: int = 3, dilations: Tuple[int] = (1, 3, 5)
+    ):
+        """
+        Initializes the ResBlock.
+
+        Args:
+            channels (int): Number of input and output channels for the convolution layers.
+            kernel_size (int): Size of the convolution kernel. Defaults to 3.
+            dilations (Tuple[int]): Tuple of dilation rates for the convolution layers in the first set.
+        """
+        super().__init__()
+        # Create convolutional layers with specified dilations and initialize weights
+        self.convs1 = self._create_convs(channels, kernel_size, dilations)
+        self.convs2 = self._create_convs(channels, kernel_size, [1] * len(dilations))
+
+    @staticmethod
+    def _create_convs(channels: int, kernel_size: int, dilations: Tuple[int]):
+        """
+        Creates a list of 1D convolutional layers with specified dilations.
+
+        Args:
+            channels (int): Number of input and output channels for the convolution layers.
+            kernel_size (int): Size of the convolution kernel.
+            dilations (Tuple[int]): Tuple of dilation rates for each convolution layer.
+        """
+        layers = torch.nn.ModuleList(
+            [create_conv1d_layer(channels, kernel_size, d) for d in dilations]
+        )
+        layers.apply(init_weights)
+        return layers
+
+    def forward(self, x: torch.Tensor, x_mask: torch.Tensor = None):
+        for conv1, conv2 in zip(self.convs1, self.convs2):
+            x_residual = x
+            x = torch.nn.functional.leaky_relu(x, LRELU_SLOPE)
+            x = apply_mask(x, x_mask)
+            x = torch.nn.functional.leaky_relu(conv1(x), LRELU_SLOPE)
+            x = apply_mask(x, x_mask)
+            x = conv2(x)
+            x = x + x_residual
+        return apply_mask(x, x_mask)
+
+    def remove_weight_norm(self):
+        for conv in chain(self.convs1, self.convs2):
+            remove_weight_norm(conv)
+
+
+class Flip(torch.nn.Module):
+    """
+    Flip module for flow-based models.
+
+    This module flips the input along the time dimension.
+    """
+
+    def forward(self, x, *args, reverse=False, **kwargs):
+        x = torch.flip(x, [1])
+        if not reverse:
+            logdet = torch.zeros(x.size(0), dtype=x.dtype, device=x.device)
+            return x, logdet
+        else:
+            return x
+
+
+class ResidualCouplingBlock(torch.nn.Module):
+    """
+    Residual Coupling Block for normalizing flow.
+
+    Args:
+        channels (int): Number of channels in the input.
+        hidden_channels (int): Number of hidden channels in the coupling layer.
+        kernel_size (int): Kernel size of the convolutional layers.
+        dilation_rate (int): Dilation rate of the convolutional layers.
+        n_layers (int): Number of layers in the coupling layer.
+        n_flows (int, optional): Number of coupling layers in the block. Defaults to 4.
+        gin_channels (int, optional): Number of channels for the global conditioning input. Defaults to 0.
+    """
+
+    def __init__(
+        self,
+        channels: int,
+        hidden_channels: int,
+        kernel_size: int,
+        dilation_rate: int,
+        n_layers: int,
+        n_flows: int = 4,
+        gin_channels: int = 0,
+    ):
+        super(ResidualCouplingBlock, self).__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 = torch.nn.ModuleList()
+        for _ in range(n_flows):
+            self.flows.append(
+                ResidualCouplingLayer(
+                    channels,
+                    hidden_channels,
+                    kernel_size,
+                    dilation_rate,
+                    n_layers,
+                    gin_channels=gin_channels,
+                    mean_only=True,
+                )
+            )
+            self.flows.append(Flip())
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        x_mask: torch.Tensor,
+        g: Optional[torch.Tensor] = None,
+        reverse: bool = 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.forward(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()
+
+    def __prepare_scriptable__(self):
+        for i in range(self.n_flows):
+            for hook in self.flows[i * 2]._forward_pre_hooks.values():
+                if (
+                    hook.__module__ == "torch.nn.utils.parametrizations.weight_norm"
+                    and hook.__class__.__name__ == "WeightNorm"
+                ):
+                    torch.nn.utils.remove_weight_norm(self.flows[i * 2])
+
+        return self
+
+
+class ResidualCouplingLayer(torch.nn.Module):
+    """
+    Residual coupling layer for flow-based models.
+
+    Args:
+        channels (int): Number of channels.
+        hidden_channels (int): Number of hidden channels.
+        kernel_size (int): Size of the convolutional kernel.
+        dilation_rate (int): Dilation rate of the convolution.
+        n_layers (int): Number of convolutional layers.
+        p_dropout (float, optional): Dropout probability. Defaults to 0.
+        gin_channels (int, optional): Number of conditioning channels. Defaults to 0.
+        mean_only (bool, optional): Whether to use mean-only coupling. Defaults to False.
+    """
+
+    def __init__(
+        self,
+        channels: int,
+        hidden_channels: int,
+        kernel_size: int,
+        dilation_rate: int,
+        n_layers: int,
+        p_dropout: float = 0,
+        gin_channels: int = 0,
+        mean_only: bool = 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 = torch.nn.Conv1d(self.half_channels, hidden_channels, 1)
+        self.enc = WaveNet(
+            hidden_channels,
+            kernel_size,
+            dilation_rate,
+            n_layers,
+            p_dropout=p_dropout,
+            gin_channels=gin_channels,
+        )
+        self.post = torch.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: torch.Tensor,
+        x_mask: torch.Tensor,
+        g: Optional[torch.Tensor] = None,
+        reverse: bool = 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()

mvsepless/vbach_lib/algorithm/synthesizers.py

@@ -0,0 +1,243 @@
+import torch
+from typing import Optional
+from .generators.hifigan_mrf import HiFiGANMRFGenerator
+from .generators.hifigan_nsf import HiFiGANNSFGenerator
+from .generators.hifigan import HiFiGANGenerator
+from .generators.refinegan import RefineGANGenerator
+from .commons import slice_segments, rand_slice_segments
+from .residuals import ResidualCouplingBlock
+from .encoders import TextEncoder, PosteriorEncoder
+
+
+class Synthesizer(torch.nn.Module):
+    """
+    Base Synthesizer model.
+
+    Args:
+        spec_channels (int): Number of channels in the spectrogram.
+        segment_size (int): Size of the audio segment.
+        inter_channels (int): Number of channels in the intermediate layers.
+        hidden_channels (int): Number of channels in the hidden layers.
+        filter_channels (int): Number of channels in the filter layers.
+        n_heads (int): Number of attention heads.
+        n_layers (int): Number of layers in the encoder.
+        kernel_size (int): Size of the convolution kernel.
+        p_dropout (float): Dropout probability.
+        resblock (str): Type of residual block.
+        resblock_kernel_sizes (list): Kernel sizes for the residual blocks.
+        resblock_dilation_sizes (list): Dilation sizes for the residual blocks.
+        upsample_rates (list): Upsampling rates for the decoder.
+        upsample_initial_channel (int): Number of channels in the initial upsampling layer.
+        upsample_kernel_sizes (list): Kernel sizes for the upsampling layers.
+        spk_embed_dim (int): Dimension of the speaker embedding.
+        gin_channels (int): Number of channels in the global conditioning vector.
+        sr (int): Sampling rate of the audio.
+        use_f0 (bool): Whether to use F0 information.
+        text_enc_hidden_dim (int): Hidden dimension for the text encoder.
+        kwargs: Additional keyword arguments.
+    """
+
+    def __init__(
+        self,
+        spec_channels: int,
+        segment_size: int,
+        inter_channels: int,
+        hidden_channels: int,
+        filter_channels: int,
+        n_heads: int,
+        n_layers: int,
+        kernel_size: int,
+        p_dropout: float,
+        resblock: str,
+        resblock_kernel_sizes: list,
+        resblock_dilation_sizes: list,
+        upsample_rates: list,
+        upsample_initial_channel: int,
+        upsample_kernel_sizes: list,
+        spk_embed_dim: int,
+        gin_channels: int,
+        sr: int,
+        use_f0: bool,
+        text_enc_hidden_dim: int = 768,
+        vocoder: str = "HiFi-GAN",
+        randomized: bool = True,
+        checkpointing: bool = False,
+        **kwargs,
+    ):
+        super().__init__()
+        self.segment_size = segment_size
+        self.use_f0 = use_f0
+        self.randomized = randomized
+
+        self.enc_p = TextEncoder(
+            inter_channels,
+            hidden_channels,
+            filter_channels,
+            n_heads,
+            n_layers,
+            kernel_size,
+            p_dropout,
+            text_enc_hidden_dim,
+            f0=use_f0,
+        )
+        print(f"Using {vocoder} vocoder")
+        if use_f0:
+            if vocoder == "MRF HiFi-GAN":
+                self.dec = HiFiGANMRFGenerator(
+                    in_channel=inter_channels,
+                    upsample_initial_channel=upsample_initial_channel,
+                    upsample_rates=upsample_rates,
+                    upsample_kernel_sizes=upsample_kernel_sizes,
+                    resblock_kernel_sizes=resblock_kernel_sizes,
+                    resblock_dilations=resblock_dilation_sizes,
+                    gin_channels=gin_channels,
+                    sample_rate=sr,
+                    harmonic_num=8,
+                    checkpointing=checkpointing,
+                )
+            elif vocoder == "RefineGAN":
+                self.dec = RefineGANGenerator(
+                    sample_rate=sr,
+                    downsample_rates=upsample_rates[::-1],
+                    upsample_rates=upsample_rates,
+                    start_channels=16,
+                    num_mels=inter_channels,
+                    checkpointing=checkpointing,
+                )
+            else:
+                self.dec = HiFiGANNSFGenerator(
+                    inter_channels,
+                    resblock_kernel_sizes,
+                    resblock_dilation_sizes,
+                    upsample_rates,
+                    upsample_initial_channel,
+                    upsample_kernel_sizes,
+                    gin_channels=gin_channels,
+                    sr=sr,
+                    checkpointing=checkpointing,
+                )
+        else:
+            if vocoder == "MRF HiFi-GAN":
+                print("MRF HiFi-GAN does not support training without pitch guidance.")
+                self.dec = None
+            elif vocoder == "RefineGAN":
+                print("RefineGAN does not support training without pitch guidance.")
+                self.dec = None
+            else:
+                self.dec = HiFiGANGenerator(
+                    inter_channels,
+                    resblock_kernel_sizes,
+                    resblock_dilation_sizes,
+                    upsample_rates,
+                    upsample_initial_channel,
+                    upsample_kernel_sizes,
+                    gin_channels=gin_channels,
+                )
+        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 = torch.nn.Embedding(spk_embed_dim, gin_channels)
+
+    def _remove_weight_norm_from(self, module):
+        for hook in module._forward_pre_hooks.values():
+            if getattr(hook, "__class__", None).__name__ == "WeightNorm":
+                torch.nn.utils.remove_weight_norm(module)
+
+    def remove_weight_norm(self):
+        for module in [self.dec, self.flow, self.enc_q]:
+            self._remove_weight_norm_from(module)
+
+    def __prepare_scriptable__(self):
+        self.remove_weight_norm()
+        return self
+
+    def forward(
+        self,
+        phone: torch.Tensor,
+        phone_lengths: torch.Tensor,
+        pitch: Optional[torch.Tensor] = None,
+        pitchf: Optional[torch.Tensor] = None,
+        y: Optional[torch.Tensor] = None,
+        y_lengths: Optional[torch.Tensor] = None,
+        ds: Optional[torch.Tensor] = None,
+    ):
+        g = self.emb_g(ds).unsqueeze(-1)
+        m_p, logs_p, x_mask = self.enc_p(phone, pitch, phone_lengths)
+
+        if y is not None:
+            z, m_q, logs_q, y_mask = self.enc_q(y, y_lengths, g=g)
+            z_p = self.flow(z, y_mask, g=g)
+            # regular old training method using random slices
+            if self.randomized:
+                z_slice, ids_slice = rand_slice_segments(
+                    z, y_lengths, self.segment_size
+                )
+                if self.use_f0:
+                    pitchf = slice_segments(pitchf, ids_slice, self.segment_size, 2)
+                    o = self.dec(z_slice, pitchf, g=g)
+                else:
+                    o = self.dec(z_slice, g=g)
+                return o, ids_slice, x_mask, y_mask, (z, z_p, m_p, logs_p, m_q, logs_q)
+            # future use for finetuning using the entire dataset each pass
+            else:
+                if self.use_f0:
+                    o = self.dec(z, pitchf, g=g)
+                else:
+                    o = self.dec(z, g=g)
+                return o, None, x_mask, y_mask, (z, z_p, m_p, logs_p, m_q, logs_q)
+        else:
+            return None, None, x_mask, None, (None, None, m_p, logs_p, None, None)
+
+    @torch.jit.export
+    def infer(
+        self,
+        phone: torch.Tensor,
+        phone_lengths: torch.Tensor,
+        pitch: Optional[torch.Tensor] = None,
+        nsff0: Optional[torch.Tensor] = None,
+        sid: torch.Tensor = None,
+        rate: Optional[torch.Tensor] = None,
+    ):
+        """
+        Inference of the model.
+
+        Args:
+            phone (torch.Tensor): Phoneme sequence.
+            phone_lengths (torch.Tensor): Lengths of the phoneme sequences.
+            pitch (torch.Tensor, optional): Pitch sequence.
+            nsff0 (torch.Tensor, optional): Fine-grained pitch sequence.
+            sid (torch.Tensor): Speaker embedding.
+            rate (torch.Tensor, optional): Rate for time-stretching.
+        """
+        g = self.emb_g(sid).unsqueeze(-1)
+        m_p, logs_p, x_mask = self.enc_p(phone, pitch, phone_lengths)
+        z_p = (m_p + torch.exp(logs_p) * torch.randn_like(m_p) * 0.66666) * x_mask
+
+        if rate is not None:
+            head = int(z_p.shape[2] * (1.0 - rate.item()))
+            z_p, x_mask = z_p[:, :, head:], x_mask[:, :, head:]
+            if self.use_f0 and nsff0 is not None:
+                nsff0 = nsff0[:, head:]
+
+        z = self.flow(z_p, x_mask, g=g, reverse=True)
+        o = (
+            self.dec(z * x_mask, nsff0, g=g)
+            if self.use_f0
+            else self.dec(z * x_mask, g=g)
+        )
+
+        return o, x_mask, (z, z_p, m_p, logs_p)