![\\\"Open](\\\"https://colab.research.google.com/assets/colab-badge.svg\\\"){kind=icon}
\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"# GPT-SoVITS Infer\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"## Env Setup (Run Once Only)\\n\",\n \"## 环境配置, 只需运行一次\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"### 1.\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": null,\n \"metadata\": {\n \"id\": \"e9b7iFV3dm1f\"\n },\n \"outputs\": [],\n \"source\": [\n \"%%writefile /content/setup.sh\\n\",\n \"set -e\\n\",\n \"\\n\",\n \"cd /content\\n\",\n \"\\n\",\n \"git clone https://github.com/RVC-Boss/GPT-SoVITS.git\\n\",\n \"\\n\",\n \"cd GPT-SoVITS\\n\",\n \"\\n\",\n \"mkdir -p GPT_weights\\n\",\n \"\\n\",\n \"mkdir -p SoVITS_weights\\n\",\n \"\\n\",\n \"if conda env list | awk '{print $1}' | grep -Fxq \\\"GPTSoVITS\\\"; then\\n\",\n \" :\\n\",\n \"else\\n\",\n \" conda create -n GPTSoVITS python=3.10 -y\\n\",\n \"fi\\n\",\n \"\\n\",\n \"source activate GPTSoVITS\\n\",\n \"\\n\",\n \"pip install ipykernel\\n\",\n \"\\n\",\n \"bash install.sh --device CU126 --source HF\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"### 2.\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": null,\n \"metadata\": {\n \"cellView\": \"form\",\n \"id\": \"0NgxXg5sjv7z\"\n },\n \"outputs\": [],\n \"source\": [\n \"%pip install -q condacolab\\n\",\n \"import condacolab\\n\",\n \"condacolab.install_from_url(\\\"https://repo.anaconda.com/archive/Anaconda3-2024.10-1-Linux-x86_64.sh\\\")\\n\",\n \"!cd /content && bash setup.sh\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"# Download Model\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"### Download From HuggingFace\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": null,\n \"metadata\": {\n \"cellView\": \"form\",\n \"id\": \"vbZY-LnM0tzq\"\n },\n \"outputs\": [],\n \"source\": [\n \"# Modify These\\n\",\n \"USER_ID = \\\"AkitoP\\\"\\n\",\n \"REPO_NAME = \\\"GPT-SoVITS-v2-aegi\\\"\\n\",\n \"BRANCH = \\\"main\\\"\\n\",\n \"GPT_PATH = \\\"new_aegigoe-e100.ckpt\\\"\\n\",\n \"SOVITS_PATH = \\\"new_aegigoe_e60_s32220.pth\\\"\\n\",\n \"\\n\",\n \"# Do Not Modify\\n\",\n \"HF_BASE = \\\"https://huggingface.co\\\"\\n\",\n \"REPO_ID = f\\\"{USER_ID}/{REPO_NAME}\\\"\\n\",\n \"GPT_URL = f\\\"{HF_BASE}/{REPO_ID}/blob/{BRANCH}/{GPT_PATH}\\\"\\n\",\n \"SOVITS_URL = f\\\"{HF_BASE}/{REPO_ID}/blob/{BRANCH}/{SOVITS_PATH}\\\"\\n\",\n \"\\n\",\n \"!cd \\\"/content/GPT-SoVITS/GPT_weights\\\" && wget \\\"{GPT_URL}\\\"\\n\",\n \"!cd \\\"/content/GPT-SoVITS/SoVITS_weights\\\" && wget \\\"{SOVITS_URL}\\\"\\n\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"### Download From ModelScope\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": null,\n \"metadata\": {},\n \"outputs\": [],\n \"source\": [\n \"# Modify These\\n\",\n \"USER_ID = \\\"aihobbyist\\\"\\n\",\n \"REPO_NAME = \\\"GPT-SoVits-V2-models\\\"\\n\",\n \"BRANCH = \\\"master\\\"\\n\",\n \"GPT_PATH = \\\"Genshin_Impact/EN/GPT_GenshinImpact_EN_5.1.ckpt\\\"\\n\",\n \"SOVITS_PATH = \\\"Wuthering_Waves/CN/SV_WutheringWaves_CN_1.3.pth\\\"\\n\",\n \"\\n\",\n \"# Do Not Modify\\n\",\n \"HF_BASE = \\\"https://www.modelscope.cn/models\\\"\\n\",\n \"REPO_ID = f\\\"{USER_ID}/{REPO_NAME}\\\"\\n\",\n \"GPT_URL = f\\\"{HF_BASE}/{REPO_ID}/resolve/{BRANCH}/{GPT_PATH}\\\"\\n\",\n \"SOVITS_URL = f\\\"{HF_BASE}/{REPO_ID}/resolve/{BRANCH}/{SOVITS_PATH}\\\"\\n\",\n \"\\n\",\n \"!cd \\\"/content/GPT-SoVITS/GPT_weights\\\" && wget \\\"{GPT_URL}\\\"\\n\",\n \"!cd \\\"/content/GPT-SoVITS/SoVITS_weights\\\" && wget \\\"{SOVITS_URL}\\\"\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"# Launch WebUI\\n\",\n \"# 启动 WebUI\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": null,\n \"metadata\": {\n \"cellView\": \"form\",\n \"id\": \"4oRGUzkrk8C7\"\n },\n \"outputs\": [],\n \"source\": [\n \"!cd /content/GPT-SoVITS && source activate GPTSoVITS && export is_share=True && python webui.py\"\n ]\n }\n ],\n \"metadata\": {\n \"accelerator\": \"GPU\",\n \"colab\": {\n \"provenance\": []\n },\n \"kernelspec\": {\n \"display_name\": \"Python 3\",\n \"name\": \"python3\"\n }\n },\n \"nbformat\": 4,\n \"nbformat_minor\": 0\n}\n"
},
{
"path": "Colab-WebUI.ipynb",
"content": "{\n \"cells\": [\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {\n \"colab_type\": \"text\",\n \"id\": \"view-in-github\"\n },\n \"source\": [\n \"\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"# GPT-SoVITS WebUI\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {\n \"id\": \"_o6a8GS2lWQM\"\n },\n \"source\": [\n \"## Env Setup (Run Once Only)\\n\",\n \"## 环境配置, 只需运行一次\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"### 1.\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": null,\n \"metadata\": {},\n \"outputs\": [],\n \"source\": [\n \"%%writefile /content/setup.sh\\n\",\n \"set -e\\n\",\n \"\\n\",\n \"cd /content\\n\",\n \"\\n\",\n \"git clone https://github.com/RVC-Boss/GPT-SoVITS.git\\n\",\n \"\\n\",\n \"cd GPT-SoVITS\\n\",\n \"\\n\",\n \"if conda env list | awk '{print $1}' | grep -Fxq \\\"GPTSoVITS\\\"; then\\n\",\n \" :\\n\",\n \"else\\n\",\n \" conda create -n GPTSoVITS python=3.10 -y\\n\",\n \"fi\\n\",\n \"\\n\",\n \"source activate GPTSoVITS\\n\",\n \"\\n\",\n \"pip install ipykernel\\n\",\n \"\\n\",\n \"bash install.sh --device CU126 --source HF --download-uvr5\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"### 2.\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": null,\n \"metadata\": {},\n \"outputs\": [],\n \"source\": [\n \"%pip install -q condacolab\\n\",\n \"import condacolab\\n\",\n \"condacolab.install_from_url(\\\"https://repo.anaconda.com/archive/Anaconda3-2024.10-1-Linux-x86_64.sh\\\")\\n\",\n \"!cd /content && bash setup.sh\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"## Launch WebUI\\n\",\n \"## 启动 WebUI\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": null,\n \"metadata\": {\n \"id\": \"4oRGUzkrk8C7\"\n },\n \"outputs\": [],\n \"source\": [\n \"!cd /content/GPT-SoVITS && source activate GPTSoVITS && export is_share=True && python webui.py\"\n ]\n }\n ],\n \"metadata\": {\n \"accelerator\": \"GPU\",\n \"colab\": {\n \"include_colab_link\": true,\n \"provenance\": []\n },\n \"kernelspec\": {\n \"display_name\": \"Python 3\",\n \"name\": \"python3\"\n }\n },\n \"nbformat\": 4,\n \"nbformat_minor\": 0\n}\n"
},
{
"path": "Docker/install_wrapper.sh",
"content": "#!/bin/bash\n\nSCRIPT_DIR=\"$(cd \"$(dirname \"${BASH_SOURCE[0]}\")\" &>/dev/null && pwd)\"\n\ncd \"$SCRIPT_DIR\" || exit 1\n\ncd .. || exit 1\n\nset -e\n\nsource \"$HOME/conda/etc/profile.d/conda.sh\"\n\nmkdir -p GPT_SoVITS\n\nmkdir -p GPT_SoVITS/text\n\nln -s /workspace/models/pretrained_models /workspace/GPT-SoVITS/GPT_SoVITS/pretrained_models\n\nln -s /workspace/models/G2PWModel /workspace/GPT-SoVITS/GPT_SoVITS/text/G2PWModel\n\nTERM=dumb bash install.sh --device \"CU${CUDA_VERSION//./}\" --source HF\n\npip cache purge\n\npip show torch\n\nrm -rf /tmp/* /var/tmp/*\n\nrm -rf \"$HOME/conda/pkgs\"\n\nmkdir -p \"$HOME/conda/pkgs\"\n\nrm -rf /root/.conda /root/.cache\n"
},
{
"path": "Docker/miniforge_install.sh",
"content": "#!/bin/bash\n\nset -e\n\nSCRIPT_DIR=\"$(cd \"$(dirname \"${BASH_SOURCE[0]}\")\" &>/dev/null && pwd)\"\n\ncd \"$SCRIPT_DIR\" || exit 1\n\ncd .. || exit 1\n\nif [ -d \"$HOME/conda\" ]; then\n exit 0\nfi\n\nWORKFLOW=${WORKFLOW:-\"false\"}\nTARGETPLATFORM=${TARGETPLATFORM:-\"linux/amd64\"}\n\nif [ \"$WORKFLOW\" = \"true\" ]; then\n WGET_CMD=(wget -nv --tries=25 --wait=5 --read-timeout=40 --retry-on-http-error=404)\nelse\n WGET_CMD=(wget --tries=25 --wait=5 --read-timeout=40 --retry-on-http-error=404)\nfi\n\nif [ \"$TARGETPLATFORM\" = \"linux/amd64\" ]; then\n \"${WGET_CMD[@]}\" -O Miniforge.sh \"https://github.com/conda-forge/miniforge/releases/latest/download/Miniforge3-$(uname)-$(uname -m).sh\"\n SYSROOT_PKG=\"sysroot_linux-64>=2.28\"\nelif [ \"$TARGETPLATFORM\" = \"linux/arm64\" ]; then\n \"${WGET_CMD[@]}\" -O Miniforge.sh \"https://github.com/conda-forge/miniforge/releases/latest/download/Miniforge3-$(uname)-$(uname -m).sh\"\n SYSROOT_PKG=\"sysroot_linux-aarch64>=2.28\"\nelse\n exit 1\nfi\n\nLOG_PATH=\"/tmp/miniforge-install.log\"\n\nbash Miniforge.sh -b -p \"$HOME/conda\" >\"$LOG_PATH\" 2>&1\n\nif [ $? -eq 0 ]; then\n echo \"== Miniforge Installed ==\"\nelse\n echo \"Failed to Install miniforge\"\n tail -n 50 \"$LOG_PATH\"\n exit 1\nfi\n\nrm Miniforge.sh\n\nsource \"$HOME/conda/etc/profile.d/conda.sh\"\n\n\"$HOME/conda/bin/conda\" init bash\n\nsource \"$HOME/.bashrc\"\n\n\"$HOME/conda/bin/conda\" info\n\n\"$HOME/conda/bin/conda\" update --all -y\n\n\"$HOME/conda/bin/conda\" install python=3.12 -y\n\n\"$HOME/conda/bin/conda\" install gcc=11 gxx ffmpeg cmake make unzip $SYSROOT_PKG \"libstdcxx-ng>=11\" -y\n\nif [ \"$CUDA_VERSION\" = \"12.8\" ]; then\n \"$HOME/conda/bin/pip\" install torch torchcodec --no-cache-dir --index-url https://download.pytorch.org/whl/cu128\n \"$HOME/conda/bin/conda\" install cuda-nvcc=12.8 -y\nelif [ \"$CUDA_VERSION\" = \"12.6\" ]; then\n \"$HOME/conda/bin/pip\" install torch torchcodec --no-cache-dir --index-url https://download.pytorch.org/whl/cu126\n \"$HOME/conda/bin/conda\" install cuda-nvcc=12.6 -y\nfi\n\nexport PATH=\"$HOME/conda/bin:$PATH\"\n\n\"$HOME/conda/bin/pip\" install psutil ninja packaging wheel \"setuptools>=42\" einops\n\"$HOME/conda/bin/pip\" install flash-attn -i https://xxxxrt666.github.io/PIP-Index/ --no-build-isolation\n\"$HOME/conda/bin/pip\" cache purge\n\nrm $LOG_PATH\n\nrm -rf \"$HOME/conda/pkgs\"\n\nmkdir -p \"$HOME/conda/pkgs\"\n\nrm -rf \"$HOME/.conda\" \"$HOME/.cache\""
},
{
"path": "Dockerfile",
"content": "ARG CUDA_VERSION=12.6\nARG TORCH_BASE=full\n\nFROM xxxxrt666/torch-base:cu${CUDA_VERSION}-${TORCH_BASE}\n\nLABEL maintainer=\"XXXXRT\"\nLABEL version=\"V2 Pro\"\nLABEL description=\"Docker image for GPT-SoVITS\"\n\nARG CUDA_VERSION=12.6\n\nENV CUDA_VERSION=${CUDA_VERSION}\n\nSHELL [\"/bin/bash\", \"-c\"]\n\nWORKDIR /workspace/GPT-SoVITS\n\nCOPY Docker /workspace/GPT-SoVITS/Docker/\n\nARG LITE=false\nENV LITE=${LITE}\n\nARG WORKFLOW=false\nENV WORKFLOW=${WORKFLOW}\n\nARG TARGETPLATFORM\nENV TARGETPLATFORM=${TARGETPLATFORM}\n\nCOPY extra-req.txt /workspace/GPT-SoVITS/\n\nCOPY requirements.txt /workspace/GPT-SoVITS/\n\nCOPY install.sh /workspace/GPT-SoVITS/\n\nRUN bash Docker/install_wrapper.sh\n\nEXPOSE 9871 9872 9873 9874 9880\n\nENV PYTHONPATH=\"/workspace/GPT-SoVITS\"\n\nRUN conda init bash && echo \"conda activate base\" >> ~/.bashrc\n\nWORKDIR /workspace\n\nRUN rm -rf /workspace/GPT-SoVITS\n\nWORKDIR /workspace/GPT-SoVITS\n\nCOPY . /workspace/GPT-SoVITS\n\nCMD [\"/bin/bash\", \"-c\", \"\\\n rm -rf /workspace/GPT-SoVITS/GPT_SoVITS/pretrained_models && \\\n rm -rf /workspace/GPT-SoVITS/GPT_SoVITS/text/G2PWModel && \\\n rm -rf /workspace/GPT-SoVITS/tools/asr/models && \\\n rm -rf /workspace/GPT-SoVITS/tools/uvr5/uvr5_weights && \\\n ln -s /workspace/models/pretrained_models /workspace/GPT-SoVITS/GPT_SoVITS/pretrained_models && \\\n ln -s /workspace/models/G2PWModel /workspace/GPT-SoVITS/GPT_SoVITS/text/G2PWModel && \\\n ln -s /workspace/models/asr_models /workspace/GPT-SoVITS/tools/asr/models && \\\n ln -s /workspace/models/uvr5_weights /workspace/GPT-SoVITS/tools/uvr5/uvr5_weights && \\\n exec bash\"]"
},
{
"path": "GPT_SoVITS/AR/__init__.py",
"content": ""
},
{
"path": "GPT_SoVITS/AR/data/__init__.py",
"content": ""
},
{
"path": "GPT_SoVITS/AR/data/bucket_sampler.py",
"content": "# modified from https://github.com/yangdongchao/SoundStorm/blob/master/soundstorm/s1/AR/data/bucket_sampler.py\n# reference: https://github.com/lifeiteng/vall-e\nimport itertools\nimport math\nimport random\nfrom random import shuffle\nfrom typing import Iterator, Optional, TypeVar\n\nimport torch\nimport torch.distributed as dist\nfrom torch.utils.data import Dataset, Sampler\n\n__all__ = [\n \"DistributedBucketSampler\",\n]\n\nT_co = TypeVar(\"T_co\", covariant=True)\n\n\nclass DistributedBucketSampler(Sampler[T_co]):\n r\"\"\"\n sort the dataset wrt. input length\n divide samples into buckets\n sort within buckets\n divide buckets into batches\n sort batches\n \"\"\"\n\n def __init__(\n self,\n dataset: Dataset,\n num_replicas: Optional[int] = None,\n rank: Optional[int] = None,\n shuffle: bool = True,\n seed: int = 0,\n drop_last: bool = False,\n batch_size: int = 32,\n ) -> None:\n if num_replicas is None:\n if not dist.is_available():\n raise RuntimeError(\"Requires distributed package to be available\")\n num_replicas = dist.get_world_size() if torch.cuda.is_available() else 1\n if rank is None:\n if not dist.is_available():\n raise RuntimeError(\"Requires distributed package to be available\")\n rank = dist.get_rank() if torch.cuda.is_available() else 0\n if torch.cuda.is_available():\n torch.cuda.set_device(rank)\n if rank >= num_replicas or rank < 0:\n raise ValueError(\"Invalid rank {}, rank should be in the interval [0, {}]\".format(rank, num_replicas - 1))\n self.dataset = dataset\n self.num_replicas = num_replicas\n self.rank = rank\n self.epoch = 0\n self.drop_last = drop_last\n # If the dataset length is evenly divisible by # of replicas, then there\n # is no need to drop any data, since the dataset will be split equally.\n if self.drop_last and len(self.dataset) % self.num_replicas != 0: # type: ignore[arg-type]\n # Split to nearest available length that is evenly divisible.\n # This is to ensure each rank receives the same amount of data when\n # using this Sampler.\n self.num_samples = math.ceil(\n (len(self.dataset) - self.num_replicas) / self.num_replicas, # type: ignore[arg-type]\n )\n else:\n self.num_samples = math.ceil(\n len(self.dataset) / self.num_replicas,\n ) # type: ignore[arg-type]\n self.total_size = self.num_samples * self.num_replicas\n self.shuffle = shuffle\n self.seed = seed\n self.batch_size = batch_size\n self.id_with_length = self._get_sample_lengths()\n self.id_buckets = self.make_buckets(bucket_width=2.0)\n\n def _get_sample_lengths(self):\n id_with_lengths = []\n for i in range(len(self.dataset)):\n id_with_lengths.append((i, self.dataset.get_sample_length(i)))\n id_with_lengths.sort(key=lambda x: x[1])\n return id_with_lengths\n\n def make_buckets(self, bucket_width: float = 2.0):\n buckets = []\n cur = []\n max_sec = bucket_width\n for id, sec in self.id_with_length:\n if sec < max_sec:\n cur.append(id)\n else:\n buckets.append(cur)\n cur = [id]\n max_sec += bucket_width\n if len(cur) > 0:\n buckets.append(cur)\n return buckets\n\n def __iter__(self) -> Iterator[T_co]:\n if self.shuffle:\n # deterministically shuffle based on epoch and seed\n g = torch.Generator()\n g.manual_seed(self.seed + self.epoch)\n random.seed(self.epoch + self.seed)\n shuffled_bucket = []\n for buc in self.id_buckets:\n buc_copy = buc.copy()\n shuffle(buc_copy)\n shuffled_bucket.append(buc_copy)\n grouped_batch_size = self.batch_size * self.num_replicas\n shuffled_bucket = list(itertools.chain(*shuffled_bucket))\n n_batch = int(math.ceil(len(shuffled_bucket) / grouped_batch_size))\n batches = [shuffled_bucket[b * grouped_batch_size : (b + 1) * grouped_batch_size] for b in range(n_batch)]\n shuffle(batches)\n indices = list(itertools.chain(*batches))\n else:\n # type: ignore[arg-type]\n indices = list(range(len(self.dataset)))\n\n if not self.drop_last:\n # add extra samples to make it evenly divisible\n padding_size = self.total_size - len(indices)\n if padding_size <= len(indices):\n indices += indices[:padding_size]\n else:\n indices += (indices * math.ceil(padding_size / len(indices)))[:padding_size]\n else:\n # remove tail of data to make it evenly divisible.\n indices = indices[: self.total_size]\n assert len(indices) == self.total_size\n\n # subsample\n indices = indices[self.rank : self.total_size : self.num_replicas]\n assert len(indices) == self.num_samples\n\n return iter(indices)\n\n def __len__(self) -> int:\n return self.num_samples\n\n def set_epoch(self, epoch: int) -> None:\n r\"\"\"\n Sets the epoch for this sampler. When :attr:`shuffle=True`, this ensures all replicas\n use a different random ordering for each epoch. Otherwise, the next iteration of this\n sampler will yield the same ordering.\n\n Args:\n epoch (int): Epoch number.\n \"\"\"\n self.epoch = epoch\n"
},
{
"path": "GPT_SoVITS/AR/data/data_module.py",
"content": "# modified from https://github.com/yangdongchao/SoundStorm/blob/master/soundstorm/s1/AR/data/data_module.py\n# reference: https://github.com/lifeiteng/vall-e\nfrom pytorch_lightning import LightningDataModule\nfrom torch.utils.data import DataLoader\n\nfrom AR.data.bucket_sampler import DistributedBucketSampler\nfrom AR.data.dataset import Text2SemanticDataset\n\n\nclass Text2SemanticDataModule(LightningDataModule):\n def __init__(\n self,\n config,\n train_semantic_path,\n train_phoneme_path,\n dev_semantic_path=None,\n dev_phoneme_path=None,\n ):\n super().__init__()\n self.config = config\n self.train_semantic_path = train_semantic_path\n self.train_phoneme_path = train_phoneme_path\n self.dev_semantic_path = dev_semantic_path\n self.dev_phoneme_path = dev_phoneme_path\n self.num_workers = self.config[\"data\"][\"num_workers\"]\n\n def prepare_data(self):\n pass\n\n def setup(self, stage=None, output_logs=False):\n self._train_dataset = Text2SemanticDataset(\n phoneme_path=self.train_phoneme_path,\n semantic_path=self.train_semantic_path,\n max_sec=self.config[\"data\"][\"max_sec\"],\n pad_val=self.config[\"data\"][\"pad_val\"],\n )\n self._dev_dataset = self._train_dataset\n # self._dev_dataset = Text2SemanticDataset(\n # phoneme_path=self.dev_phoneme_path,\n # semantic_path=self.dev_semantic_path,\n # max_sample=self.config['data']['max_eval_sample'],\n # max_sec=self.config['data']['max_sec'],\n # pad_val=self.config['data']['pad_val'])\n\n def train_dataloader(self):\n batch_size = (\n self.config[\"train\"][\"batch_size\"] // 2\n if self.config[\"train\"].get(\"if_dpo\", False) is True\n else self.config[\"train\"][\"batch_size\"]\n )\n batch_size = max(min(batch_size, len(self._train_dataset) // 4), 1) # 防止不保存\n sampler = DistributedBucketSampler(self._train_dataset, batch_size=batch_size)\n return DataLoader(\n self._train_dataset,\n batch_size=batch_size,\n sampler=sampler,\n collate_fn=self._train_dataset.collate,\n num_workers=self.num_workers,\n persistent_workers=True,\n prefetch_factor=16,\n )\n\n def val_dataloader(self):\n return DataLoader(\n self._dev_dataset,\n batch_size=1,\n shuffle=False,\n collate_fn=self._train_dataset.collate,\n num_workers=max(self.num_workers, 12),\n persistent_workers=True,\n prefetch_factor=16,\n )\n\n # 这个会使用到嘛？\n def test_dataloader(self):\n return DataLoader(\n self._dev_dataset,\n batch_size=1,\n shuffle=False,\n collate_fn=self._train_dataset.collate,\n )\n"
},
{
"path": "GPT_SoVITS/AR/data/dataset.py",
"content": "# modified from https://github.com/yangdongchao/SoundStorm/blob/master/soundstorm/s1/AR/data/dataset.py\n# reference: https://github.com/lifeiteng/vall-e\n\n# sys.path.append(\"/data/docker/liujing04/gpt-vits/mq-vits-s1bert_no_bert\")\nimport os\nimport traceback\nfrom typing import Dict, List\n\nimport numpy as np\nimport pandas as pd\nimport torch\nfrom torch.utils.data import DataLoader, Dataset\n\nversion = os.environ.get(\"version\", None)\n\nfrom text import cleaned_text_to_sequence\n\n# from config import exp_dir\n\n\ndef batch_sequences(sequences: List[np.array], axis: int = 0, pad_value: int = 0):\n seq = sequences[0]\n ndim = seq.ndim\n if axis < 0:\n axis += ndim\n dtype = seq.dtype\n pad_value = dtype.type(pad_value)\n seq_lengths = [seq.shape[axis] for seq in sequences]\n max_length = np.max(seq_lengths)\n\n padded_sequences = []\n for seq, length in zip(sequences, seq_lengths):\n padding = [(0, 0)] * axis + [(0, max_length - length)] + [(0, 0)] * (ndim - axis - 1)\n padded_seq = np.pad(seq, padding, mode=\"constant\", constant_values=pad_value)\n padded_sequences.append(padded_seq)\n batch = np.stack(padded_sequences)\n return batch\n\n\nclass Text2SemanticDataset(Dataset):\n \"\"\"dataset class for text tokens to semantic model training.\"\"\"\n\n def __init__(\n self,\n phoneme_path: str,\n semantic_path: str,\n max_sample: int = None,\n max_sec: int = 100,\n pad_val: int = 1024,\n # min value of phoneme/sec\n min_ps_ratio: int = 3,\n # max value of phoneme/sec\n max_ps_ratio: int = 25,\n ) -> None:\n super().__init__()\n\n self.semantic_data = pd.read_csv(\n semantic_path,\n delimiter=\"\\t\",\n encoding=\"utf-8\",\n )\n # get dict\n self.path2 = phoneme_path # \"%s/2-name2text.txt\"%exp_dir#phoneme_path\n self.path3 = \"%s/3-bert\" % (\n os.path.dirname(\n phoneme_path,\n )\n ) # \"%s/3-bert\"%exp_dir#bert_dir\n self.path6 = semantic_path # \"%s/6-name2semantic.tsv\"%exp_dir#semantic_path\n assert os.path.exists(self.path2)\n assert os.path.exists(self.path6)\n self.phoneme_data = {}\n with open(self.path2, \"r\", encoding=\"utf8\") as f:\n lines = f.read().strip(\"\\n\").split(\"\\n\")\n\n for line in lines:\n tmp = line.split(\"\\t\")\n if len(tmp) != 4:\n continue\n self.phoneme_data[tmp[0]] = [tmp[1], tmp[2], tmp[3]]\n\n # self.phoneme_data = np.load(phoneme_path, allow_pickle=True).item()\n # pad for semantic tokens\n self.PAD: int = pad_val\n # self.hz = 25\n # with open(\"/data/docker/liujing04/gpt-vits/mq-vits-s1bert_no_bert/configs/s2.json\", \"r\") as f:data = f.read()\n # data=json.loads(data)[\"model\"][\"semantic_frame_rate\"]#50hz\n # self.hz=int(data[:-2])#\n self.hz = int(os.environ.get(\"hz\", \"25hz\")[:-2])\n\n # max seconds of semantic token\n self.max_sec = max_sec\n self.min_ps_ratio = min_ps_ratio\n self.max_ps_ratio = max_ps_ratio\n\n if max_sample is not None:\n self.semantic_data = self.semantic_data[:max_sample]\n\n # {idx: (semantic, phoneme)}\n # semantic list, phoneme list\n self.semantic_phoneme = []\n self.item_names = []\n\n self.inited = False\n\n if not self.inited:\n # 调用初始化函数\n self.init_batch()\n self.inited = True\n del self.semantic_data\n del self.phoneme_data\n # self.tokenizer = AutoTokenizer.from_pretrained(\"hfl/chinese-roberta-wwm-ext-large\")\n # self.tokenizer = AutoTokenizer.from_pretrained(\"/data/docker/liujing04/bert-vits2/Bert-VITS2-master20231106/bert/chinese-roberta-wwm-ext-large\")\n\n def init_batch(self):\n semantic_data_len = len(self.semantic_data)\n phoneme_data_len = len(self.phoneme_data.keys())\n print(\"semantic_data_len:\", semantic_data_len)\n print(\"phoneme_data_len:\", phoneme_data_len)\n print(self.semantic_data)\n idx = 0\n num_not_in = 0\n num_deleted_bigger = 0\n num_deleted_ps = 0\n for i in range(semantic_data_len):\n # 先依次遍历\n # get str\n item_name = self.semantic_data.iloc[i, 0]\n # print(self.phoneme_data)\n try:\n phoneme, word2ph, text = self.phoneme_data[item_name]\n except Exception:\n traceback.print_exc()\n # print(f\"{item_name} not in self.phoneme_data !\")\n num_not_in += 1\n continue\n\n semantic_str = self.semantic_data.iloc[i, 1]\n # get token list\n semantic_ids = [int(idx) for idx in semantic_str.split(\" \")]\n # (T), 是否需要变成 (1, T) -> 不需要，因为需要求 len\n # 过滤掉太长的样本\n if (\n len(semantic_ids) > self.max_sec * self.hz\n ): #########1###根据token个数推测总时长过滤时长60s（config里）#40*25=1k\n num_deleted_bigger += 1\n continue\n # (T, ), 这个速度不会很慢，所以可以在一开始就处理，无需在 __getitem__ 里面单个处理####\n phoneme = phoneme.split(\" \")\n\n try:\n phoneme_ids = cleaned_text_to_sequence(phoneme, version)\n except:\n traceback.print_exc()\n # print(f\"{item_name} not in self.phoneme_data !\")\n num_not_in += 1\n continue\n # if len(phoneme_ids) >400:###########2：改为恒定限制为semantic/2.5就行\n if len(phoneme_ids) > self.max_sec * self.hz / 2.5: ###########2：改为恒定限制为semantic/2.5就行\n num_deleted_ps += 1\n continue\n # if len(semantic_ids) > 1000:###########3\n # num_deleted_bigger += 1\n # continue\n\n ps_ratio = len(phoneme_ids) / (len(semantic_ids) / self.hz)\n\n if ps_ratio > self.max_ps_ratio or ps_ratio < self.min_ps_ratio: ##########4#3~25#每秒多少个phone\n num_deleted_ps += 1\n # print(item_name)\n continue\n\n self.semantic_phoneme.append((semantic_ids, phoneme_ids))\n idx += 1\n self.item_names.append(item_name)\n\n min_num = 100 # 20直接不补#30补了也不存ckpt\n leng = len(self.semantic_phoneme)\n if leng < min_num:\n tmp1 = self.semantic_phoneme\n tmp2 = self.item_names\n self.semantic_phoneme = []\n self.item_names = []\n for _ in range(max(2, int(min_num / leng))):\n self.semantic_phoneme += tmp1\n self.item_names += tmp2\n if num_not_in > 0:\n print(f\"there are {num_not_in} semantic datas not in phoneme datas\")\n if num_deleted_bigger > 0:\n print(\n f\"deleted {num_deleted_bigger} audios who's duration are bigger than {self.max_sec} seconds\",\n )\n if num_deleted_ps > 0:\n # 4702 for LibriTTS, LirbriTTS 是标注数据, 是否需要筛？=> 需要，有值为 100 的极端值\n print(\n f\"deleted {num_deleted_ps} audios who's phoneme/sec are bigger than {self.max_ps_ratio} or smaller than {self.min_ps_ratio}\",\n )\n \"\"\"\n there are 31 semantic datas not in phoneme datas\n deleted 34 audios who's duration are bigger than 54 seconds\n deleted 3190 audios who's phoneme/sec are bigger than 25 or smaller than 3\n dataset.__len__(): 366463\n\n \"\"\"\n # 345410 for LibriTTS\n print(\"dataset.__len__():\", self.__len__())\n\n def __get_item_names__(self) -> List[str]:\n return self.item_names\n\n def __len__(self) -> int:\n return len(self.semantic_phoneme)\n\n def __getitem__(self, idx: int) -> Dict:\n semantic_ids, phoneme_ids = self.semantic_phoneme[idx]\n item_name = self.item_names[idx]\n phoneme_ids_len = len(phoneme_ids)\n # semantic tokens target\n semantic_ids_len = len(semantic_ids)\n\n flag = 0\n path_bert = \"%s/%s.pt\" % (self.path3, item_name)\n if os.path.exists(path_bert) == True:\n bert_feature = torch.load(path_bert, map_location=\"cpu\")\n else:\n flag = 1\n if flag == 1:\n # bert_feature=torch.zeros_like(phoneme_ids,dtype=torch.float32)\n bert_feature = None\n else:\n assert bert_feature.shape[-1] == len(phoneme_ids)\n return {\n \"idx\": idx,\n \"phoneme_ids\": phoneme_ids,\n \"phoneme_ids_len\": phoneme_ids_len,\n \"semantic_ids\": semantic_ids,\n \"semantic_ids_len\": semantic_ids_len,\n \"bert_feature\": bert_feature,\n }\n\n def get_sample_length(self, idx: int):\n semantic_ids = self.semantic_phoneme[idx][0]\n sec = 1.0 * len(semantic_ids) / self.hz\n return sec\n\n def collate(self, examples: List[Dict]) -> Dict:\n sample_index: List[int] = []\n phoneme_ids: List[torch.Tensor] = []\n phoneme_ids_lens: List[int] = []\n semantic_ids: List[torch.Tensor] = []\n semantic_ids_lens: List[int] = []\n # return\n\n for item in examples:\n sample_index.append(item[\"idx\"])\n phoneme_ids.append(np.array(item[\"phoneme_ids\"], dtype=np.int64))\n semantic_ids.append(np.array(item[\"semantic_ids\"], dtype=np.int64))\n phoneme_ids_lens.append(item[\"phoneme_ids_len\"])\n semantic_ids_lens.append(item[\"semantic_ids_len\"])\n\n # pad 0\n phoneme_ids = batch_sequences(phoneme_ids)\n semantic_ids = batch_sequences(semantic_ids, pad_value=self.PAD)\n\n # # convert each batch to torch.tensor\n phoneme_ids = torch.tensor(phoneme_ids)\n semantic_ids = torch.tensor(semantic_ids)\n phoneme_ids_lens = torch.tensor(phoneme_ids_lens)\n semantic_ids_lens = torch.tensor(semantic_ids_lens)\n bert_padded = torch.FloatTensor(len(examples), 1024, max(phoneme_ids_lens))\n bert_padded.zero_()\n\n for idx, item in enumerate(examples):\n bert = item[\"bert_feature\"]\n if bert != None:\n bert_padded[idx, :, : bert.shape[-1]] = bert\n\n return {\n # List[int]\n \"ids\": sample_index,\n # torch.Tensor (B, max_phoneme_length)\n \"phoneme_ids\": phoneme_ids,\n # torch.Tensor (B)\n \"phoneme_ids_len\": phoneme_ids_lens,\n # torch.Tensor (B, max_semantic_ids_length)\n \"semantic_ids\": semantic_ids,\n # torch.Tensor (B)\n \"semantic_ids_len\": semantic_ids_lens,\n # torch.Tensor (B, 1024, max_phoneme_length)\n \"bert_feature\": bert_padded,\n }\n\n\nif __name__ == \"__main__\":\n root_dir = \"/data/docker/liujing04/gpt-vits/prepare/dump_mix/\"\n dataset = Text2SemanticDataset(\n phoneme_path=root_dir + \"phoneme_train.npy\",\n semantic_path=root_dir + \"semantic_train.tsv\",\n )\n\n batch_size = 12\n dataloader = DataLoader(\n dataset,\n batch_size=batch_size,\n collate_fn=dataset.collate,\n shuffle=False,\n )\n for i, batch in enumerate(dataloader):\n if i % 1000 == 0:\n print(i)\n # if i == 0:\n # print('batch[\"ids\"]:', batch[\"ids\"])\n # print('batch[\"phoneme_ids\"]:', batch[\"phoneme_ids\"],\n # batch[\"phoneme_ids\"].shape)\n # print('batch[\"phoneme_ids_len\"]:', batch[\"phoneme_ids_len\"],\n # batch[\"phoneme_ids_len\"].shape)\n # print('batch[\"semantic_ids\"]:', batch[\"semantic_ids\"],\n # batch[\"semantic_ids\"].shape)\n # print('batch[\"semantic_ids_len\"]:', batch[\"semantic_ids_len\"],\n # batch[\"semantic_ids_len\"].shape)\n"
},
{
"path": "GPT_SoVITS/AR/models/__init__.py",
"content": ""
},
{
"path": "GPT_SoVITS/AR/models/t2s_lightning_module.py",
"content": "# modified from https://github.com/yangdongchao/SoundStorm/blob/master/soundstorm/s1/AR/models/t2s_lightning_module.py\n# reference: https://github.com/lifeiteng/vall-e\nimport os\nimport sys\n\nnow_dir = os.getcwd()\nsys.path.append(now_dir)\nfrom typing import Dict\n\nimport torch\nfrom pytorch_lightning import LightningModule\n\nfrom AR.models.t2s_model import Text2SemanticDecoder\nfrom AR.modules.lr_schedulers import WarmupCosineLRSchedule\nfrom AR.modules.optim import ScaledAdam\n\n\nclass Text2SemanticLightningModule(LightningModule):\n def __init__(self, config, output_dir, is_train=True):\n super().__init__()\n self.config = config\n self.top_k = 3\n self.model = Text2SemanticDecoder(config=config, top_k=self.top_k)\n pretrained_s1 = config.get(\"pretrained_s1\")\n if pretrained_s1 and is_train:\n # print(self.load_state_dict(torch.load(pretrained_s1,map_location=\"cpu\")[\"state_dict\"]))\n print(\n self.load_state_dict(\n torch.load(\n pretrained_s1,\n map_location=\"cpu\",\n weights_only=False,\n )[\"weight\"],\n )\n )\n if is_train:\n self.automatic_optimization = False\n self.save_hyperparameters()\n self.eval_dir = output_dir / \"eval\"\n self.eval_dir.mkdir(parents=True, exist_ok=True)\n\n def training_step(self, batch: Dict, batch_idx: int):\n opt = self.optimizers()\n scheduler = self.lr_schedulers()\n forward = self.model.forward if self.config[\"train\"].get(\"if_dpo\", False) == True else self.model.forward_old\n loss, acc = forward(\n batch[\"phoneme_ids\"],\n batch[\"phoneme_ids_len\"],\n batch[\"semantic_ids\"],\n batch[\"semantic_ids_len\"],\n batch[\"bert_feature\"],\n )\n self.manual_backward(loss)\n if batch_idx > 0 and batch_idx % 4 == 0:\n opt.step()\n opt.zero_grad()\n scheduler.step()\n\n self.log(\n \"total_loss\",\n loss,\n on_step=True,\n on_epoch=True,\n prog_bar=True,\n sync_dist=True,\n )\n self.log(\n \"lr\",\n scheduler.get_last_lr()[0],\n on_epoch=True,\n prog_bar=True,\n sync_dist=True,\n )\n self.log(\n f\"top_{self.top_k}_acc\",\n acc,\n on_step=True,\n on_epoch=True,\n prog_bar=True,\n sync_dist=True,\n )\n\n def validation_step(self, batch: Dict, batch_idx: int):\n return\n\n # # get loss\n # loss, acc = self.model.forward(\n # batch['phoneme_ids'], batch['phoneme_ids_len'],\n # batch['semantic_ids'], batch['semantic_ids_len'],\n # batch['bert_feature']\n # )\n #\n # self.log(\n # \"val_total_loss\",\n # loss,\n # on_step=True,\n # on_epoch=True,\n # prog_bar=True,\n # sync_dist=True)\n # self.log(\n # f\"val_top_{self.top_k}_acc\",\n # acc,\n # on_step=True,\n # on_epoch=True,\n # prog_bar=True,\n # sync_dist=True)\n #\n # # get infer output\n # semantic_len = batch['semantic_ids'].size(1)\n # prompt_len = min(int(semantic_len * 0.5), 150)\n # prompt = batch['semantic_ids'][:, :prompt_len]\n # pred_semantic = self.model.infer(batch['phoneme_ids'],\n # batch['phoneme_ids_len'], prompt,\n # batch['bert_feature']\n # )\n # save_name = f'semantic_toks_{batch_idx}.pt'\n # save_path = os.path.join(self.eval_dir, save_name)\n # torch.save(pred_semantic.detach().cpu(), save_path)\n\n def configure_optimizers(self):\n model_parameters = self.model.parameters()\n parameters_names = []\n parameters_names.append([name_param_pair[0] for name_param_pair in self.model.named_parameters()])\n lm_opt = ScaledAdam(\n model_parameters,\n lr=0.01,\n betas=(0.9, 0.95),\n clipping_scale=2.0,\n parameters_names=parameters_names,\n show_dominant_parameters=False,\n clipping_update_period=1000,\n )\n\n return {\n \"optimizer\": lm_opt,\n \"lr_scheduler\": {\n \"scheduler\": WarmupCosineLRSchedule(\n lm_opt,\n init_lr=self.config[\"optimizer\"][\"lr_init\"],\n peak_lr=self.config[\"optimizer\"][\"lr\"],\n end_lr=self.config[\"optimizer\"][\"lr_end\"],\n warmup_steps=self.config[\"optimizer\"][\"warmup_steps\"],\n total_steps=self.config[\"optimizer\"][\"decay_steps\"],\n )\n },\n }\n"
},
{
"path": "GPT_SoVITS/AR/models/t2s_lightning_module_onnx.py",
"content": "# modified from https://github.com/yangdongchao/SoundStorm/blob/master/soundstorm/s1/AR/models/t2s_lightning_module.py\n# reference: https://github.com/lifeiteng/vall-e\nimport os\nimport sys\n\nnow_dir = os.getcwd()\nsys.path.append(now_dir)\nfrom typing import Dict\n\nimport torch\nfrom pytorch_lightning import LightningModule\n\nfrom AR.models.t2s_model_onnx import Text2SemanticDecoder\nfrom AR.modules.lr_schedulers import WarmupCosineLRSchedule\nfrom AR.modules.optim import ScaledAdam\n\n\nclass Text2SemanticLightningModule(LightningModule):\n def __init__(self, config, output_dir, is_train=True):\n super().__init__()\n self.config = config\n self.top_k = 3\n self.model = Text2SemanticDecoder(config=config, top_k=self.top_k)\n pretrained_s1 = config.get(\"pretrained_s1\")\n if pretrained_s1 and is_train:\n # print(self.load_state_dict(torch.load(pretrained_s1,map_location=\"cpu\")[\"state_dict\"]))\n print(\n self.load_state_dict(\n torch.load(\n pretrained_s1,\n map_location=\"cpu\",\n )[\"weight\"],\n ),\n )\n if is_train:\n self.automatic_optimization = False\n self.save_hyperparameters()\n self.eval_dir = output_dir / \"eval\"\n self.eval_dir.mkdir(parents=True, exist_ok=True)\n\n def training_step(self, batch: Dict, batch_idx: int):\n opt = self.optimizers()\n scheduler = self.lr_schedulers()\n loss, acc = self.model.forward(\n batch[\"phoneme_ids\"],\n batch[\"phoneme_ids_len\"],\n batch[\"semantic_ids\"],\n batch[\"semantic_ids_len\"],\n batch[\"bert_feature\"],\n )\n self.manual_backward(loss)\n if batch_idx > 0 and batch_idx % 4 == 0:\n opt.step()\n opt.zero_grad()\n scheduler.step()\n\n self.log(\n \"total_loss\",\n loss,\n on_step=True,\n on_epoch=True,\n prog_bar=True,\n sync_dist=True,\n )\n self.log(\n \"lr\",\n scheduler.get_last_lr()[0],\n on_epoch=True,\n prog_bar=True,\n sync_dist=True,\n )\n self.log(\n f\"top_{self.top_k}_acc\",\n acc,\n on_step=True,\n on_epoch=True,\n prog_bar=True,\n sync_dist=True,\n )\n\n def validation_step(self, batch: Dict, batch_idx: int):\n return\n\n def configure_optimizers(self):\n model_parameters = self.model.parameters()\n parameters_names = []\n parameters_names.append([name_param_pair[0] for name_param_pair in self.model.named_parameters()])\n lm_opt = ScaledAdam(\n model_parameters,\n lr=0.01,\n betas=(0.9, 0.95),\n clipping_scale=2.0,\n parameters_names=parameters_names,\n show_dominant_parameters=False,\n clipping_update_period=1000,\n )\n\n return {\n \"optimizer\": lm_opt,\n \"lr_scheduler\": {\n \"scheduler\": WarmupCosineLRSchedule(\n lm_opt,\n init_lr=self.config[\"optimizer\"][\"lr_init\"],\n peak_lr=self.config[\"optimizer\"][\"lr\"],\n end_lr=self.config[\"optimizer\"][\"lr_end\"],\n warmup_steps=self.config[\"optimizer\"][\"warmup_steps\"],\n total_steps=self.config[\"optimizer\"][\"decay_steps\"],\n )\n },\n }\n"
},
{
"path": "GPT_SoVITS/AR/models/t2s_model.py",
"content": "# modified from https://github.com/yangdongchao/SoundStorm/blob/master/soundstorm/s1/AR/models/t2s_model.py\n# reference: https://github.com/lifeiteng/vall-e\nimport math\nfrom typing import List, Optional\n\nimport torch\nfrom torch import nn\nfrom torch.nn import functional as F\nfrom torchmetrics.classification import MulticlassAccuracy\nfrom tqdm import tqdm\n\nfrom AR.models.utils import (\n dpo_loss,\n get_batch_logps,\n make_pad_mask,\n make_pad_mask_left,\n make_reject_y,\n sample,\n topk_sampling,\n)\nfrom AR.modules.embedding import SinePositionalEmbedding, TokenEmbedding\nfrom AR.modules.transformer import LayerNorm, TransformerEncoder, TransformerEncoderLayer\n\ndefault_config = {\n \"embedding_dim\": 512,\n \"hidden_dim\": 512,\n \"num_head\": 8,\n \"num_layers\": 12,\n \"num_codebook\": 8,\n \"p_dropout\": 0.0,\n \"vocab_size\": 1024 + 1,\n \"phoneme_vocab_size\": 512,\n \"EOS\": 1024,\n}\n\n\n# @torch.jit.script ## 使用的话首次推理会非常慢，而且推理速度不稳定\n# Efficient implementation equivalent to the following:\ndef scaled_dot_product_attention(\n query: torch.Tensor,\n key: torch.Tensor,\n value: torch.Tensor,\n attn_mask: Optional[torch.Tensor] = None,\n scale: Optional[torch.Tensor] = None,\n) -> torch.Tensor:\n B, H, L, S = query.size(0), query.size(1), query.size(-2), key.size(-2)\n if scale is None:\n scale_factor = torch.tensor(1 / math.sqrt(query.size(-1)))\n else:\n scale_factor = scale\n attn_bias = torch.zeros(B, H, L, S, dtype=query.dtype, device=query.device)\n\n if attn_mask is not None:\n if attn_mask.dtype == torch.bool:\n attn_bias.masked_fill_(attn_mask, float(\"-inf\"))\n else:\n attn_bias += attn_mask\n attn_weight = query @ key.transpose(-2, -1) * scale_factor\n attn_weight += attn_bias\n attn_weight = torch.softmax(attn_weight, dim=-1)\n\n if attn_mask is not None:\n if attn_mask.dtype == torch.bool:\n attn_weight.masked_fill_(attn_mask, 0)\n else:\n attn_mask[attn_mask != float(\"-inf\")] = 0\n attn_mask[attn_mask == float(\"-inf\")] = 1\n attn_weight.masked_fill_(attn_mask, 0)\n\n return attn_weight @ value\n\n\n@torch.jit.script\nclass T2SMLP:\n def __init__(self, w1, b1, w2, b2):\n self.w1 = w1\n self.b1 = b1\n self.w2 = w2\n self.b2 = b2\n\n def forward(self, x):\n x = F.relu(F.linear(x, self.w1, self.b1))\n x = F.linear(x, self.w2, self.b2)\n return x\n\n\n@torch.jit.script\nclass T2SBlock:\n def __init__(\n self,\n num_heads,\n hidden_dim: int,\n mlp: T2SMLP,\n qkv_w,\n qkv_b,\n out_w,\n out_b,\n norm_w1,\n norm_b1,\n norm_eps1,\n norm_w2,\n norm_b2,\n norm_eps2,\n ):\n self.num_heads = num_heads\n self.mlp = mlp\n self.hidden_dim: int = hidden_dim\n self.qkv_w = qkv_w\n self.qkv_b = qkv_b\n self.out_w = out_w\n self.out_b = out_b\n self.norm_w1 = norm_w1\n self.norm_b1 = norm_b1\n self.norm_eps1 = norm_eps1\n self.norm_w2 = norm_w2\n self.norm_b2 = norm_b2\n self.norm_eps2 = norm_eps2\n\n self.false = torch.tensor(False, dtype=torch.bool)\n\n @torch.jit.ignore\n def to_mask(\n self,\n x: torch.Tensor,\n padding_mask: Optional[torch.Tensor],\n ):\n if padding_mask is None:\n return x\n\n if padding_mask.dtype == torch.bool:\n return x.masked_fill(padding_mask, 0)\n else:\n return x * padding_mask\n\n def process_prompt(\n self,\n x: torch.Tensor,\n attn_mask: torch.Tensor,\n padding_mask: Optional[torch.Tensor] = None,\n torch_sdpa: bool = True,\n ):\n q, k, v = F.linear(self.to_mask(x, padding_mask), self.qkv_w, self.qkv_b).chunk(3, dim=-1)\n\n batch_size = q.shape[0]\n q_len = q.shape[1]\n kv_len = k.shape[1]\n\n q = self.to_mask(q, padding_mask)\n k_cache = self.to_mask(k, padding_mask)\n v_cache = self.to_mask(v, padding_mask)\n\n q = q.view(batch_size, q_len, self.num_heads, -1).transpose(1, 2)\n k = k_cache.view(batch_size, kv_len, self.num_heads, -1).transpose(1, 2)\n v = v_cache.view(batch_size, kv_len, self.num_heads, -1).transpose(1, 2)\n\n if torch_sdpa:\n attn = F.scaled_dot_product_attention(q, k, v, ~attn_mask)\n else:\n attn = scaled_dot_product_attention(q, k, v, attn_mask)\n\n attn = attn.transpose(1, 2).reshape(batch_size, q_len, -1)\n attn = F.linear(self.to_mask(attn, padding_mask), self.out_w, self.out_b)\n\n x = x + attn\n x = F.layer_norm(x, [self.hidden_dim], self.norm_w1, self.norm_b1, self.norm_eps1)\n x = x + self.mlp.forward(x)\n x = F.layer_norm(\n x,\n [self.hidden_dim],\n self.norm_w2,\n self.norm_b2,\n self.norm_eps2,\n )\n return x, k_cache, v_cache\n\n def decode_next_token(\n self,\n x: torch.Tensor,\n k_cache: torch.Tensor,\n v_cache: torch.Tensor,\n attn_mask: torch.Tensor = None,\n torch_sdpa: bool = True,\n ):\n q, k, v = F.linear(x, self.qkv_w, self.qkv_b).chunk(3, dim=-1)\n\n k_cache = torch.cat([k_cache, k], dim=1)\n v_cache = torch.cat([v_cache, v], dim=1)\n\n batch_size = q.shape[0]\n q_len = q.shape[1]\n kv_len = k_cache.shape[1]\n\n q = q.view(batch_size, q_len, self.num_heads, -1).transpose(1, 2)\n k = k_cache.view(batch_size, kv_len, self.num_heads, -1).transpose(1, 2)\n v = v_cache.view(batch_size, kv_len, self.num_heads, -1).transpose(1, 2)\n\n if torch_sdpa:\n attn = F.scaled_dot_product_attention(q, k, v, (~attn_mask) if attn_mask is not None else None)\n else:\n attn = scaled_dot_product_attention(q, k, v, attn_mask)\n\n attn = attn.transpose(1, 2).reshape(batch_size, q_len, -1)\n attn = F.linear(attn, self.out_w, self.out_b)\n\n x = x + attn\n x = F.layer_norm(\n x,\n [self.hidden_dim],\n self.norm_w1,\n self.norm_b1,\n self.norm_eps1,\n )\n x = x + self.mlp.forward(x)\n x = F.layer_norm(\n x,\n [self.hidden_dim],\n self.norm_w2,\n self.norm_b2,\n self.norm_eps2,\n )\n return x, k_cache, v_cache\n\n\n@torch.jit.script\nclass T2STransformer:\n def __init__(self, num_blocks: int, blocks: List[T2SBlock]):\n self.num_blocks: int = num_blocks\n self.blocks = blocks\n\n def process_prompt(\n self,\n x: torch.Tensor,\n attn_mask: torch.Tensor,\n padding_mask: Optional[torch.Tensor] = None,\n torch_sdpa: bool = True,\n ):\n k_cache: List[torch.Tensor] = []\n v_cache: List[torch.Tensor] = []\n for i in range(self.num_blocks):\n x, k_cache_, v_cache_ = self.blocks[i].process_prompt(x, attn_mask, padding_mask, torch_sdpa)\n k_cache.append(k_cache_)\n v_cache.append(v_cache_)\n return x, k_cache, v_cache\n\n def decode_next_token(\n self,\n x: torch.Tensor,\n k_cache: List[torch.Tensor],\n v_cache: List[torch.Tensor],\n attn_mask: torch.Tensor = None,\n torch_sdpa: bool = True,\n ):\n for i in range(self.num_blocks):\n x, k_cache[i], v_cache[i] = self.blocks[i].decode_next_token(\n x, k_cache[i], v_cache[i], attn_mask, torch_sdpa\n )\n return x, k_cache, v_cache\n\n\nclass Text2SemanticDecoder(nn.Module):\n def __init__(self, config, norm_first=False, top_k=3):\n super(Text2SemanticDecoder, self).__init__()\n self.model_dim = config[\"model\"][\"hidden_dim\"]\n self.embedding_dim = config[\"model\"][\"embedding_dim\"]\n self.num_head = config[\"model\"][\"head\"]\n self.num_layers = config[\"model\"][\"n_layer\"]\n self.norm_first = norm_first\n self.vocab_size = config[\"model\"][\"vocab_size\"]\n self.phoneme_vocab_size = config[\"model\"][\"phoneme_vocab_size\"]\n self.p_dropout = config[\"model\"][\"dropout\"]\n self.EOS = config[\"model\"][\"EOS\"]\n self.norm_first = norm_first\n assert self.EOS == self.vocab_size - 1\n # should be same as num of kmeans bin\n # assert self.EOS == 1024\n self.bert_proj = nn.Linear(1024, self.embedding_dim)\n self.ar_text_embedding = TokenEmbedding(\n self.embedding_dim,\n self.phoneme_vocab_size,\n self.p_dropout,\n )\n self.ar_text_position = SinePositionalEmbedding(\n self.embedding_dim,\n dropout=0.1,\n scale=False,\n alpha=True,\n )\n self.ar_audio_embedding = TokenEmbedding(\n self.embedding_dim,\n self.vocab_size,\n self.p_dropout,\n )\n self.ar_audio_position = SinePositionalEmbedding(\n self.embedding_dim,\n dropout=0.1,\n scale=False,\n alpha=True,\n )\n\n self.h = TransformerEncoder(\n TransformerEncoderLayer(\n d_model=self.model_dim,\n nhead=self.num_head,\n dim_feedforward=self.model_dim * 4,\n dropout=0.1,\n batch_first=True,\n norm_first=norm_first,\n ),\n num_layers=self.num_layers,\n norm=LayerNorm(self.model_dim) if norm_first else None,\n )\n\n self.ar_predict_layer = nn.Linear(self.model_dim, self.vocab_size, bias=False)\n self.loss_fct = nn.CrossEntropyLoss(reduction=\"sum\")\n\n self.ar_accuracy_metric = MulticlassAccuracy(\n self.vocab_size,\n top_k=top_k,\n average=\"micro\",\n multidim_average=\"global\",\n ignore_index=self.EOS,\n )\n\n blocks = []\n\n for i in range(self.num_layers):\n layer = self.h.layers[i]\n t2smlp = T2SMLP(\n layer.linear1.weight,\n layer.linear1.bias,\n layer.linear2.weight,\n layer.linear2.bias,\n )\n\n block = T2SBlock(\n self.num_head,\n self.model_dim,\n t2smlp,\n layer.self_attn.in_proj_weight,\n layer.self_attn.in_proj_bias,\n layer.self_attn.out_proj.weight,\n layer.self_attn.out_proj.bias,\n layer.norm1.weight,\n layer.norm1.bias,\n layer.norm1.eps,\n layer.norm2.weight,\n layer.norm2.bias,\n layer.norm2.eps,\n )\n\n blocks.append(block)\n\n self.t2s_transformer = T2STransformer(self.num_layers, blocks)\n\n def make_input_data(self, x, x_lens, y, y_lens, bert_feature):\n x = self.ar_text_embedding(x)\n x = x + self.bert_proj(bert_feature.transpose(1, 2))\n x = self.ar_text_position(x)\n x_mask = make_pad_mask_left(x_lens)\n\n y_mask = make_pad_mask(y_lens)\n y_mask_int = y_mask.type(torch.int64)\n codes = y.type(torch.int64) * (1 - y_mask_int)\n\n # Training\n # AR Decoder\n y, targets = self.pad_y_eos(codes, y_mask_int, eos_id=self.EOS)\n x_len = x_lens.max()\n y_len = y_lens.max()\n y_emb = self.ar_audio_embedding(y)\n y_pos = self.ar_audio_position(y_emb)\n\n xy_padding_mask = torch.concat([x_mask, y_mask], dim=1)\n\n ar_xy_padding_mask = xy_padding_mask\n\n x_attn_mask = F.pad(\n torch.zeros((x_len, x_len), dtype=torch.bool, device=x.device),\n (0, y_len),\n value=True,\n )\n # x_attn_mask[:, x_len]=False\n y_attn_mask = F.pad(\n torch.triu(\n torch.ones(y_len, y_len, dtype=torch.bool, device=x.device),\n diagonal=1,\n ),\n (x_len, 0),\n value=False,\n )\n\n xy_attn_mask = torch.concat([x_attn_mask, y_attn_mask], dim=0)\n bsz, src_len = x.shape[0], x_len + y_len\n _xy_padding_mask = (\n ar_xy_padding_mask.view(bsz, 1, 1, src_len)\n .expand(-1, self.num_head, -1, -1)\n .reshape(bsz * self.num_head, 1, src_len)\n )\n xy_attn_mask = xy_attn_mask.logical_or(_xy_padding_mask)\n new_attn_mask = torch.zeros_like(xy_attn_mask, dtype=x.dtype)\n new_attn_mask.masked_fill_(xy_attn_mask, float(\"-inf\"))\n xy_attn_mask = new_attn_mask\n # x 和完整的 y 一次性输入模型\n xy_pos = torch.concat([x, y_pos], dim=1)\n\n return xy_pos, xy_attn_mask, targets\n\n def forward(self, x, x_lens, y, y_lens, bert_feature):\n \"\"\"\n x: phoneme_ids\n y: semantic_ids\n \"\"\"\n\n reject_y, reject_y_lens = make_reject_y(y, y_lens)\n\n xy_pos, xy_attn_mask, targets = self.make_input_data(x, x_lens, y, y_lens, bert_feature)\n\n xy_dec, _ = self.h(\n (xy_pos, None),\n mask=xy_attn_mask,\n )\n x_len = x_lens.max()\n logits = self.ar_predict_layer(xy_dec[:, x_len-1:])\n\n ###### DPO #############\n reject_xy_pos, reject_xy_attn_mask, reject_targets = self.make_input_data(\n x, x_lens, reject_y, reject_y_lens, bert_feature\n )\n\n reject_xy_dec, _ = self.h(\n (reject_xy_pos, None),\n mask=reject_xy_attn_mask,\n )\n x_len = x_lens.max()\n reject_logits = self.ar_predict_layer(reject_xy_dec[:, x_len-1:])\n\n # loss\n # from feiteng: 每次 duration 越多, 梯度更新也应该更多, 所以用 sum\n\n loss_1 = F.cross_entropy(logits.permute(0, 2, 1), targets, reduction=\"sum\")\n acc = self.ar_accuracy_metric(logits.permute(0, 2, 1).detach(), targets).item()\n\n A_logits, R_logits = get_batch_logps(logits, reject_logits, targets, reject_targets)\n loss_2, _, _ = dpo_loss(A_logits, R_logits, 0, 0, 0.2, reference_free=True)\n\n loss = loss_1 + loss_2\n\n return loss, acc\n\n def forward_old(self, x, x_lens, y, y_lens, bert_feature):\n \"\"\"\n x: phoneme_ids\n y: semantic_ids\n \"\"\"\n x = self.ar_text_embedding(x)\n x = x + self.bert_proj(bert_feature.transpose(1, 2))\n x = self.ar_text_position(x)\n x_mask = make_pad_mask_left(x_lens)\n\n y_mask = make_pad_mask(y_lens)\n y_mask_int = y_mask.type(torch.int64)\n codes = y.type(torch.int64) * (1 - y_mask_int)\n\n # Training\n # AR Decoder\n y, targets = self.pad_y_eos(codes, y_mask_int, eos_id=self.EOS)\n x_len = x_lens.max()\n y_len = y_lens.max()\n y_emb = self.ar_audio_embedding(y)\n y_pos = self.ar_audio_position(y_emb)\n\n xy_padding_mask = torch.concat([x_mask, y_mask], dim=1)\n ar_xy_padding_mask = xy_padding_mask\n\n x_attn_mask = F.pad(\n torch.zeros((x_len, x_len), dtype=torch.bool, device=x.device),\n (0, y_len),\n value=True,\n )\n y_attn_mask = F.pad(\n torch.triu(\n torch.ones(y_len, y_len, dtype=torch.bool, device=x.device),\n diagonal=1,\n ),\n (x_len, 0),\n value=False,\n )\n xy_attn_mask = torch.concat([x_attn_mask, y_attn_mask], dim=0)\n bsz, src_len = x.shape[0], x_len + y_len\n _xy_padding_mask = (\n ar_xy_padding_mask.view(bsz, 1, 1, src_len)\n .expand(-1, self.num_head, -1, -1)\n .reshape(bsz * self.num_head, 1, src_len)\n )\n xy_attn_mask = xy_attn_mask.logical_or(_xy_padding_mask)\n new_attn_mask = torch.zeros_like(xy_attn_mask, dtype=x.dtype)\n new_attn_mask.masked_fill_(xy_attn_mask, float(\"-inf\"))\n xy_attn_mask = new_attn_mask\n # x 和完整的 y 一次性输入模型\n xy_pos = torch.concat([x, y_pos], dim=1)\n xy_dec, _ = self.h(\n (xy_pos, None),\n mask=xy_attn_mask,\n )\n logits = self.ar_predict_layer(xy_dec[:, x_len-1:]).permute(0, 2, 1)\n # loss\n # from feiteng: 每次 duration 越多, 梯度更新也应该更多, 所以用 sum\n loss = F.cross_entropy(logits, targets, reduction=\"sum\")\n acc = self.ar_accuracy_metric(logits.detach(), targets).item()\n return loss, acc\n\n # 需要看下这个函数和 forward 的区别以及没有 semantic 的时候 prompts 输入什么\n def infer(\n self,\n x,\n x_lens,\n prompts,\n bert_feature,\n top_k: int = -100,\n early_stop_num: int = -1,\n temperature: float = 1.0,\n ):\n x = self.ar_text_embedding(x)\n x = x + self.bert_proj(bert_feature.transpose(1, 2))\n x = self.ar_text_position(x)\n\n # AR Decoder\n y = prompts\n prefix_len = y.shape[1]\n x_len = x.shape[1]\n x_attn_mask = torch.zeros((x_len, x_len), dtype=torch.bool)\n stop = False\n for _ in tqdm(range(1500)):\n y_emb = self.ar_audio_embedding(y)\n y_pos = self.ar_audio_position(y_emb)\n # x 和逐渐增长的 y 一起输入给模型\n xy_pos = torch.concat([x, y_pos], dim=1)\n y_len = y.shape[1]\n x_attn_mask_pad = F.pad(\n x_attn_mask,\n (0, y_len),\n value=True,\n )\n y_attn_mask = F.pad(\n torch.triu(torch.ones(y_len, y_len, dtype=torch.bool), diagonal=1),\n (x_len, 0),\n value=False,\n )\n xy_attn_mask = torch.concat([x_attn_mask_pad, y_attn_mask], dim=0).to(y.device)\n\n xy_dec, _ = self.h(\n (xy_pos, None),\n mask=xy_attn_mask,\n )\n logits = self.ar_predict_layer(xy_dec[:, -1])\n samples = topk_sampling(logits, top_k=top_k, top_p=1.0, temperature=temperature)\n\n if early_stop_num != -1 and (y.shape[1] - prefix_len) > early_stop_num:\n print(\"use early stop num:\", early_stop_num)\n stop = True\n\n if torch.argmax(logits, dim=-1)[0] == self.EOS or samples[0, 0] == self.EOS:\n # print(torch.argmax(logits, dim=-1)[0] == self.EOS, samples[0, 0] == self.EOS)\n stop = True\n if stop:\n if prompts.shape[1] == y.shape[1]:\n y = torch.concat([y, torch.zeros_like(samples)], dim=1)\n print(\"bad zero prediction\")\n print(f\"T2S Decoding EOS [{prefix_len} -> {y.shape[1]}]\")\n break\n # 本次生成的 semantic_ids 和之前的 y 构成新的 y\n # print(samples.shape)#[1,1]#第一个1是bs\n # import os\n # os._exit(2333)\n y = torch.concat([y, samples], dim=1)\n return y\n\n def pad_y_eos(self, y, y_mask_int, eos_id):\n targets = F.pad(y, (0, 1), value=0) + eos_id * F.pad(y_mask_int, (0, 1), value=1)\n # 错位\n return targets[:, :-1], targets\n\n def infer_panel_batch_infer(\n self,\n x: List[torch.LongTensor], #####全部文本token\n x_lens: torch.LongTensor,\n prompts: torch.LongTensor, ####参考音频token\n bert_feature: List[torch.LongTensor],\n top_k: int = -100,\n top_p: int = 100,\n early_stop_num: int = -1,\n temperature: float = 1.0,\n repetition_penalty: float = 1.35,\n **kwargs,\n ):\n if prompts is None:\n print(\"Warning: Prompt free is not supported batch_infer! switch to naive_infer\")\n return self.infer_panel_naive_batched(\n x,\n x_lens,\n prompts,\n bert_feature,\n top_k=top_k,\n top_p=top_p,\n early_stop_num=early_stop_num,\n temperature=temperature,\n **kwargs,\n )\n\n max_len = kwargs.get(\"max_len\", x_lens.max())\n x_list = []\n for x_item, bert_item in zip(x, bert_feature):\n # max_len = max(max_len, x_item.shape[0], bert_item.shape[1])\n x_item = self.ar_text_embedding(x_item.unsqueeze(0))\n x_item = x_item + self.bert_proj(bert_item.transpose(0, 1).unsqueeze(0))\n x_item = self.ar_text_position(x_item).squeeze(0)\n # x_item = F.pad(x_item,(0,0,0,max_len-x_item.shape[0]),value=0) if x_item.shape[0]\n for p in group[\"params\"]:\n state = self.state[p]\n ...\n \n you can do:\n \n with self.batched_params(group[\"params\"]) as batches:\n for p, state, p_names in batches:\n ...\n \n\n Args:\n group: a parameter group, which is a list of parameters; should be\n one of self.param_groups.\n group_params_names: name for each parameter in group,\n which is List[str].\n \"\"\"\n batches = defaultdict(list) # `batches` maps from tuple (dtype_as_str,*shape) to list of nn.Parameter\n batches_names = defaultdict(list) # `batches` maps from tuple (dtype_as_str,*shape) to list of str\n\n assert len(param_group) == len(group_params_names)\n for p, named_p in zip(param_group, group_params_names):\n key = (str(p.dtype), *p.shape)\n batches[key].append(p)\n batches_names[key].append(named_p)\n\n batches_names_keys = list(batches_names.keys())\n sorted_idx = sorted(range(len(batches_names)), key=lambda i: batches_names_keys[i])\n batches_names = [batches_names[batches_names_keys[idx]] for idx in sorted_idx]\n batches = [batches[batches_names_keys[idx]] for idx in sorted_idx]\n\n stacked_params_dict = dict()\n\n # turn batches into a list, in deterministic order.\n # tuples will contain tuples of (stacked_param, state, stacked_params_names),\n # one for each batch in `batches`.\n tuples = []\n\n for batch, batch_names in zip(batches, batches_names):\n p = batch[0]\n # we arbitrarily store the state in the\n # state corresponding to the 1st parameter in the\n # group. class Optimizer will take care of saving/loading state.\n state = self.state[p]\n p_stacked = torch.stack(batch)\n grad = torch.stack([torch.zeros_like(p) if p.grad is None else p.grad for p in batch])\n p_stacked.grad = grad\n stacked_params_dict[key] = p_stacked\n tuples.append((p_stacked, state, batch_names))\n\n yield tuples # <-- calling code will do the actual optimization here!\n\n for (stacked_params, _state, _names), batch in zip(tuples, batches):\n for i, p in enumerate(batch): # batch is list of Parameter\n p.copy_(stacked_params[i])\n\n\nclass ScaledAdam(BatchedOptimizer):\n \"\"\"\n Implements 'Scaled Adam', a variant of Adam where we scale each parameter's update\n proportional to the norm of that parameter; and also learn the scale of the parameter,\n in log space, subject to upper and lower limits (as if we had factored each parameter as\n param = underlying_param * log_scale.exp())\n\n\n Args:\n params: The parameters or param_groups to optimize (like other Optimizer subclasses)\n lr: The learning rate. We will typically use a learning rate schedule that starts\n at 0.03 and decreases over time, i.e. much higher than other common\n optimizers.\n clipping_scale: (e.g. 2.0)\n A scale for gradient-clipping: if specified, the normalized gradients\n over the whole model will be clipped to have 2-norm equal to\n `clipping_scale` times the median 2-norm over the most recent period\n of `clipping_update_period` minibatches. By \"normalized gradients\",\n we mean after multiplying by the rms parameter value for this tensor\n [for non-scalars]; this is appropriate because our update is scaled\n by this quantity.\n betas: beta1,beta2 are momentum constants for regular momentum, and moving sum-sq grad.\n Must satisfy 0 < beta <= beta2 < 1.\n scalar_lr_scale: A scaling factor on the learning rate, that we use to update the\n scale of each parameter tensor and scalar parameters of the mode..\n If each parameter were decomposed\n as p * p_scale.exp(), where (p**2).mean().sqrt() == 1.0, scalar_lr_scale\n would be a the scaling factor on the learning rate of p_scale.\n eps: A general-purpose epsilon to prevent division by zero\n param_min_rms: Minimum root-mean-square value of parameter tensor, for purposes of\n learning the scale on the parameters (we'll constrain the rms of each non-scalar\n parameter tensor to be >= this value)\n param_max_rms: Maximum root-mean-square value of parameter tensor, for purposes of\n learning the scale on the parameters (we'll constrain the rms of each non-scalar\n parameter tensor to be <= this value)\n scalar_max: Maximum absolute value for scalar parameters (applicable if your\n model has any parameters with numel() == 1).\n size_update_period: The periodicity, in steps, with which we update the size (scale)\n of the parameter tensor. This is provided to save a little time\n in the update.\n clipping_update_period: if clipping_scale is specified, this is the period\n \"\"\"\n\n def __init__(\n self,\n params,\n lr=3e-02,\n clipping_scale=None,\n betas=(0.9, 0.98),\n scalar_lr_scale=0.1,\n eps=1.0e-08,\n param_min_rms=1.0e-05,\n param_max_rms=3.0,\n scalar_max=10.0,\n size_update_period=4,\n clipping_update_period=100,\n parameters_names=None,\n show_dominant_parameters=True,\n ):\n assert parameters_names is not None, (\n \"Please prepare parameters_names,which is a List[List[str]]. Each List[str] is for a groupand each str is for a parameter\"\n )\n defaults = dict(\n lr=lr,\n clipping_scale=clipping_scale,\n betas=betas,\n scalar_lr_scale=scalar_lr_scale,\n eps=eps,\n param_min_rms=param_min_rms,\n param_max_rms=param_max_rms,\n scalar_max=scalar_max,\n size_update_period=size_update_period,\n clipping_update_period=clipping_update_period,\n )\n\n super(ScaledAdam, self).__init__(params, defaults)\n assert len(self.param_groups) == len(parameters_names)\n self.parameters_names = parameters_names\n self.show_dominant_parameters = show_dominant_parameters\n\n def __setstate__(self, state):\n super(ScaledAdam, self).__setstate__(state)\n\n @torch.no_grad()\n def step(self, closure=None):\n \"\"\"Performs a single optimization step.\n\n Arguments:\n closure (callable, optional): A closure that reevaluates the model\n and returns the loss.\n \"\"\"\n loss = None\n if closure is not None:\n with torch.enable_grad():\n loss = closure()\n\n batch = True\n\n for group, group_params_names in zip(self.param_groups, self.parameters_names):\n with self.batched_params(group[\"params\"], group_params_names) as batches:\n # batches is list of pairs (stacked_param, state). stacked_param is like\n # a regular parameter, and will have a .grad, but the 1st dim corresponds to\n # a stacking dim, it is not a real dim.\n\n if len(batches[0][1]) == 0: # if len(first state) == 0: not yet initialized\n clipping_scale = 1\n else:\n clipping_scale = self._get_clipping_scale(group, batches)\n\n for p, state, _ in batches:\n # Perform optimization step.\n # grad is not going to be None, we handled that when creating the batches.\n grad = p.grad\n if grad.is_sparse:\n raise RuntimeError(\"ScaledAdam optimizer does not support sparse gradients\")\n # State initialization\n if len(state) == 0:\n self._init_state(group, p, state)\n\n self._step_one_batch(group, p, state, clipping_scale)\n\n return loss\n\n def _init_state(self, group: dict, p: Tensor, state: dict):\n \"\"\"\n Initializes state dict for parameter 'p'. Assumes that dim 0 of tensor p\n is actually the batch dimension, corresponding to batched-together\n parameters of a given shape.\n\n\n Args:\n group: Dict to look up configuration values.\n p: The parameter that we are initializing the state for\n state: Dict from string to whatever state we are initializing\n \"\"\"\n size_update_period = group[\"size_update_period\"]\n\n state[\"step\"] = 0\n\n kwargs = {\"device\": p.device, \"dtype\": p.dtype}\n\n # 'delta' implements conventional momentum. There are\n # several different kinds of update going on, so rather than\n # compute \"exp_avg\" like in Adam, we store and decay a\n # parameter-change \"delta\", which combines all forms of\n # update. this is equivalent to how it's done in Adam,\n # except for the first few steps.\n state[\"delta\"] = torch.zeros_like(p, memory_format=torch.preserve_format)\n\n batch_size = p.shape[0]\n numel = p.numel() // batch_size\n numel = p.numel()\n\n if numel > 1:\n # \"param_rms\" just periodically records the scalar root-mean-square value of\n # the parameter tensor.\n # it has a shape like (batch_size, 1, 1, 1, 1)\n param_rms = (p**2).mean(dim=list(range(1, p.ndim)), keepdim=True).sqrt()\n state[\"param_rms\"] = param_rms\n\n state[\"scale_exp_avg_sq\"] = torch.zeros_like(param_rms)\n state[\"scale_grads\"] = torch.zeros(size_update_period, *param_rms.shape, **kwargs)\n\n # exp_avg_sq is the weighted sum of scaled gradients. as in Adam.\n state[\"exp_avg_sq\"] = torch.zeros_like(p, memory_format=torch.preserve_format)\n\n def _get_clipping_scale(self, group: dict, tuples: List[Tuple[Tensor, dict, List[str]]]) -> float:\n \"\"\"\n Returns a scalar factor <= 1.0 that dictates gradient clipping, i.e. we will scale the gradients\n by this amount before applying the rest of the update.\n\n Args:\n group: the parameter group, an item in self.param_groups\n tuples: a list of tuples of (param, state, param_names)\n where param is a batched set of parameters,\n with a .grad (1st dim is batch dim)\n and state is the state-dict where optimization parameters are kept.\n param_names is a List[str] while each str is name for a parameter\n in batched set of parameters \"param\".\n \"\"\"\n assert len(tuples) >= 1\n clipping_scale = group[\"clipping_scale\"]\n (first_p, first_state, _) = tuples[0]\n step = first_state[\"step\"]\n if clipping_scale is None or step == 0:\n # no clipping. return early on step == 0 because the other\n # parameters' state won't have been initialized yet.\n return 1.0\n clipping_update_period = group[\"clipping_update_period\"]\n\n tot_sumsq = torch.tensor(0.0, device=first_p.device)\n for p, state, param_names in tuples:\n grad = p.grad\n if grad.is_sparse:\n raise RuntimeError(\"ScaledAdam optimizer does not support sparse gradients\")\n if p.numel() == p.shape[0]: # a batch of scalars\n tot_sumsq += (grad**2).sum() # sum() to change shape [1] to []\n else:\n tot_sumsq += ((grad * state[\"param_rms\"]) ** 2).sum()\n\n tot_norm = tot_sumsq.sqrt()\n if \"model_norms\" not in first_state:\n first_state[\"model_norms\"] = torch.zeros(clipping_update_period, device=p.device)\n first_state[\"model_norms\"][step % clipping_update_period] = tot_norm\n\n if step % clipping_update_period == 0:\n # Print some stats.\n # We don't reach here if step == 0 because we would have returned\n # above.\n sorted_norms = first_state[\"model_norms\"].sort()[0].to(\"cpu\")\n quartiles = []\n for n in range(0, 5):\n index = min(\n clipping_update_period - 1,\n (clipping_update_period // 4) * n,\n )\n quartiles.append(sorted_norms[index].item())\n\n median = quartiles[2]\n threshold = clipping_scale * median\n first_state[\"model_norm_threshold\"] = threshold\n percent_clipped = (\n first_state[\"num_clipped\"] * 100.0 / clipping_update_period if \"num_clipped\" in first_state else 0.0\n )\n first_state[\"num_clipped\"] = 0\n quartiles = \" \".join([\"%.3e\" % x for x in quartiles])\n logging.info(\n f\"Clipping_scale={clipping_scale}, grad-norm quartiles {quartiles}, threshold={threshold:.3e}, percent-clipped={percent_clipped:.1f}\"\n )\n\n if step < clipping_update_period:\n return 1.0 # We have not yet estimated a norm to clip to.\n else:\n try:\n model_norm_threshold = first_state[\"model_norm_threshold\"]\n except KeyError:\n logging.info(\n \"Warning: model_norm_threshold not in state: possibly you changed config when restarting, adding clipping_scale option?\"\n )\n return 1.0\n ans = min(1.0, (model_norm_threshold / (tot_norm + 1.0e-20)).item())\n if ans < 1.0:\n first_state[\"num_clipped\"] += 1\n if ans < 0.1:\n logging.warning(f\"Scaling gradients by {ans}, model_norm_threshold={model_norm_threshold}\")\n if self.show_dominant_parameters:\n assert p.shape[0] == len(param_names)\n self._show_gradient_dominating_parameter(tuples, tot_sumsq)\n return ans\n\n def _show_gradient_dominating_parameter(self, tuples: List[Tuple[Tensor, dict, List[str]]], tot_sumsq: Tensor):\n \"\"\"\n Show information of parameter which dominating tot_sumsq.\n\n Args:\n tuples: a list of tuples of (param, state, param_names)\n where param is a batched set of parameters,\n with a .grad (1st dim is batch dim)\n and state is the state-dict where optimization parameters are kept.\n param_names is a List[str] while each str is name for a parameter\n in batched set of parameters \"param\".\n tot_sumsq: sumsq of all parameters. Though it's could be calculated\n from tuples, we still pass it to save some time.\n \"\"\"\n all_sumsq_orig = {}\n for p, state, batch_param_names in tuples:\n # p is a stacked batch parameters.\n batch_grad = p.grad\n if p.numel() == p.shape[0]: # a batch of scalars\n batch_sumsq_orig = batch_grad**2\n # Dummpy values used by following `zip` statement.\n batch_rms_orig = torch.ones(p.shape[0])\n else:\n batch_rms_orig = state[\"param_rms\"]\n batch_sumsq_orig = ((batch_grad * batch_rms_orig) ** 2).sum(dim=list(range(1, batch_grad.ndim)))\n\n for name, sumsq_orig, rms, grad in zip(\n batch_param_names,\n batch_sumsq_orig,\n batch_rms_orig,\n batch_grad,\n ):\n proportion_orig = sumsq_orig / tot_sumsq\n all_sumsq_orig[name] = (proportion_orig, sumsq_orig, rms, grad)\n\n assert torch.isclose(\n sum([value[0] for value in all_sumsq_orig.values()]).cpu(),\n torch.tensor(1.0),\n )\n sorted_by_proportion = {\n k: v\n for k, v in sorted(\n all_sumsq_orig.items(),\n key=lambda item: item[1][0],\n reverse=True,\n )\n }\n dominant_param_name = next(iter(sorted_by_proportion))\n (\n dominant_proportion,\n dominant_sumsq,\n dominant_rms,\n dominant_grad,\n ) = sorted_by_proportion[dominant_param_name]\n logging.info(\n f\"Parameter Dominating tot_sumsq {dominant_param_name}\"\n f\" with proportion {dominant_proportion:.2f},\"\n f\" where dominant_sumsq=(grad_sumsq*orig_rms_sq)\"\n f\"={dominant_sumsq:.3e},\"\n f\" grad_sumsq = {(dominant_grad**2).sum():.3e},\"\n f\" orig_rms_sq={(dominant_rms**2).item():.3e}\"\n )\n\n def _step_one_batch(self, group: dict, p: Tensor, state: dict, clipping_scale: float):\n \"\"\"\n Do the step for one parameter, which is actually going to be a batch of\n `real` parameters, with dim 0 as the batch dim.\n Args:\n group: dict to look up configuration values\n p: parameter to update (actually multiple parameters stacked together\n as a batch)\n state: state-dict for p, to look up the optimizer state\n \"\"\"\n lr = group[\"lr\"]\n size_update_period = group[\"size_update_period\"]\n beta1 = group[\"betas\"][0]\n\n grad = p.grad\n if clipping_scale != 1.0:\n grad = grad * clipping_scale\n step = state[\"step\"]\n delta = state[\"delta\"]\n\n delta.mul_(beta1)\n batch_size = p.shape[0]\n numel = p.numel() // batch_size\n if numel > 1:\n # Update the size/scale of p, and set param_rms\n scale_grads = state[\"scale_grads\"]\n scale_grads[step % size_update_period] = (p * grad).sum(dim=list(range(1, p.ndim)), keepdim=True)\n if step % size_update_period == size_update_period - 1:\n param_rms = state[\"param_rms\"] # shape: (batch_size, 1, 1, ..)\n param_rms.copy_((p**2).mean(dim=list(range(1, p.ndim)), keepdim=True).sqrt())\n if step > 0:\n # self._size_update() learns the overall scale on the\n # parameter, by shrinking or expanding it.\n self._size_update(group, scale_grads, p, state)\n\n if numel == 1:\n # For parameters with 1 element we just use regular Adam.\n # Updates delta.\n self._step_scalar(group, p, state)\n else:\n self._step(group, p, state)\n\n state[\"step\"] = step + 1\n\n def _size_update(\n self,\n group: dict,\n scale_grads: Tensor,\n p: Tensor,\n state: dict,\n ) -> None:\n \"\"\"\n Called only where p.numel() > 1, this updates the scale of the parameter.\n If we imagine: p = underlying_param * scale.exp(), and we are doing\n gradient descent on underlying param and on scale, this function does the update\n on `scale`.\n\n Args:\n group: dict to look up configuration values\n scale_grads: a tensor of shape (size_update_period, batch_size, 1, 1,...) containing\n grads w.r.t. the scales.\n p: The parameter to update\n state: The state-dict of p\n \"\"\"\n\n param_rms = state[\"param_rms\"]\n beta1, beta2 = group[\"betas\"]\n size_lr = group[\"lr\"] * group[\"scalar_lr_scale\"]\n param_min_rms = group[\"param_min_rms\"]\n param_max_rms = group[\"param_max_rms\"]\n eps = group[\"eps\"]\n step = state[\"step\"]\n batch_size = p.shape[0]\n\n size_update_period = scale_grads.shape[0]\n # correct beta2 for the size update period: we will have\n # faster decay at this level.\n beta2_corr = beta2**size_update_period\n\n scale_exp_avg_sq = state[\"scale_exp_avg_sq\"] # shape: (batch_size, 1, 1, ..)\n scale_exp_avg_sq.mul_(beta2_corr).add_(\n (scale_grads**2).mean(dim=0), # mean over dim `size_update_period`\n alpha=1 - beta2_corr,\n ) # shape is (batch_size, 1, 1, ...)\n\n # The 1st time we reach here is when size_step == 1.\n size_step = (step + 1) // size_update_period\n bias_correction2 = 1 - beta2_corr**size_step\n # we don't bother with bias_correction1; this will help prevent divergence\n # at the start of training.\n\n denom = scale_exp_avg_sq.sqrt() + eps\n\n scale_step = -size_lr * (bias_correction2**0.5) * scale_grads.sum(dim=0) / denom\n\n is_too_small = param_rms < param_min_rms\n is_too_large = param_rms > param_max_rms\n\n # when the param gets too small, just don't shrink it any further.\n scale_step.masked_fill_(is_too_small, 0.0)\n # when it gets too large, stop it from getting any larger.\n scale_step.masked_fill_(is_too_large, -size_lr * size_update_period)\n delta = state[\"delta\"]\n # the factor of (1-beta1) relates to momentum.\n delta.add_(p * scale_step, alpha=(1 - beta1))\n\n def _step(self, group: dict, p: Tensor, state: dict):\n \"\"\"\n This function does the core update of self.step(), in the case where the members of\n the batch have more than 1 element.\n\n Args:\n group: A dict which will be used to look up configuration values\n p: The parameter to be updated\n grad: The grad of p\n state: The state-dict corresponding to parameter p\n\n This function modifies p.\n \"\"\"\n grad = p.grad\n lr = group[\"lr\"]\n beta1, beta2 = group[\"betas\"]\n eps = group[\"eps\"]\n param_min_rms = group[\"param_min_rms\"]\n step = state[\"step\"]\n\n exp_avg_sq = state[\"exp_avg_sq\"]\n exp_avg_sq.mul_(beta2).addcmul_(grad, grad, value=(1 - beta2))\n\n this_step = state[\"step\"] - (state[\"zero_step\"] if \"zero_step\" in state else 0)\n bias_correction2 = 1 - beta2 ** (this_step + 1)\n if bias_correction2 < 0.99:\n # note: not in-place.\n exp_avg_sq = exp_avg_sq * (1.0 / bias_correction2)\n\n denom = exp_avg_sq.sqrt()\n denom += eps\n grad = grad / denom\n\n alpha = -lr * (1 - beta1) * state[\"param_rms\"].clamp(min=param_min_rms)\n\n delta = state[\"delta\"]\n delta.add_(grad * alpha)\n p.add_(delta)\n\n def _step_scalar(self, group: dict, p: Tensor, state: dict):\n \"\"\"\n A simplified form of the core update for scalar tensors, where we cannot get a good\n estimate of the parameter rms.\n \"\"\"\n beta1, beta2 = group[\"betas\"]\n scalar_max = group[\"scalar_max\"]\n eps = group[\"eps\"]\n lr = group[\"lr\"] * group[\"scalar_lr_scale\"]\n grad = p.grad\n\n exp_avg_sq = state[\"exp_avg_sq\"] # shape: (batch_size,)\n exp_avg_sq.mul_(beta2).addcmul_(grad, grad, value=1 - beta2)\n\n # bias_correction2 is like in Adam. Don't bother with bias_correction1;\n # slower update at the start will help stability anyway.\n bias_correction2 = 1 - beta2 ** (state[\"step\"] + 1)\n denom = (exp_avg_sq / bias_correction2).sqrt() + eps\n\n delta = state[\"delta\"]\n delta.add_(grad / denom, alpha=-lr * (1 - beta1))\n p.clamp_(min=-scalar_max, max=scalar_max)\n p.add_(delta)\n"
},
{
"path": "GPT_SoVITS/AR/modules/patched_mha_with_cache.py",
"content": "from torch.nn.functional import *\r\nfrom torch.nn.functional import (\r\n _mha_shape_check,\r\n _canonical_mask,\r\n _none_or_dtype,\r\n _in_projection_packed,\r\n)\r\nimport torch\r\n# Tensor = torch.Tensor\r\n# from typing import Callable, List, Optional, Tuple, Union\r\n\r\n\r\ndef multi_head_attention_forward_patched(\r\n query,\r\n key,\r\n value,\r\n embed_dim_to_check,\r\n num_heads,\r\n in_proj_weight,\r\n in_proj_bias,\r\n bias_k,\r\n bias_v,\r\n add_zero_attn,\r\n dropout_p: float,\r\n out_proj_weight,\r\n out_proj_bias,\r\n training=True,\r\n key_padding_mask=None,\r\n need_weights=True,\r\n attn_mask=None,\r\n use_separate_proj_weight=False,\r\n q_proj_weight=None,\r\n k_proj_weight=None,\r\n v_proj_weight=None,\r\n static_k=None,\r\n static_v=None,\r\n average_attn_weights=True,\r\n is_causal=False,\r\n cache=None,\r\n):\r\n r\"\"\"\r\n Args:\r\n query, key, value: map a query and a set of key-value pairs to an output.\r\n See \"Attention Is All You Need\" for more details.\r\n embed_dim_to_check: total dimension of the model.\r\n num_heads: parallel attention heads.\r\n in_proj_weight, in_proj_bias: input projection weight and bias.\r\n bias_k, bias_v: bias of the key and value sequences to be added at dim=0.\r\n add_zero_attn: add a new batch of zeros to the key and\r\n value sequences at dim=1.\r\n dropout_p: probability of an element to be zeroed.\r\n out_proj_weight, out_proj_bias: the output projection weight and bias.\r\n training: apply dropout if is ``True``.\r\n key_padding_mask: if provided, specified padding elements in the key will\r\n be ignored by the attention. This is an binary mask. When the value is True,\r\n the corresponding value on the attention layer will be filled with -inf.\r\n need_weights: output attn_output_weights.\r\n Default: `True`\r\n Note: `needs_weight` defaults to `True`, but should be set to `False`\r\n For best performance when attention weights are not nedeeded.\r\n *Setting needs_weights to `True`\r\n leads to a significant performance degradation.*\r\n attn_mask: 2D or 3D mask that prevents attention to certain positions. A 2D mask will be broadcasted for all\r\n the batches while a 3D mask allows to specify a different mask for the entries of each batch.\r\n is_causal: If specified, applies a causal mask as attention mask, and ignores\r\n attn_mask for computing scaled dot product attention.\r\n Default: ``False``.\r\n .. warning::\r\n is_causal is provides a hint that the attn_mask is the\r\n causal mask.Providing incorrect hints can result in\r\n incorrect execution, including forward and backward\r\n compatibility.\r\n use_separate_proj_weight: the function accept the proj. weights for query, key,\r\n and value in different forms. If false, in_proj_weight will be used, which is\r\n a combination of q_proj_weight, k_proj_weight, v_proj_weight.\r\n q_proj_weight, k_proj_weight, v_proj_weight, in_proj_bias: input projection weight and bias.\r\n static_k, static_v: static key and value used for attention operators.\r\n average_attn_weights: If true, indicates that the returned ``attn_weights`` should be averaged across heads.\r\n Otherwise, ``attn_weights`` are provided separately per head. Note that this flag only has an effect\r\n when ``need_weights=True.``. Default: True\r\n\r\n\r\n Shape:\r\n Inputs:\r\n - query: :math:`(L, E)` or :math:`(L, N, E)` where L is the target sequence length, N is the batch size, E is\r\n the embedding dimension.\r\n - key: :math:`(S, E)` or :math:`(S, N, E)`, where S is the source sequence length, N is the batch size, E is\r\n the embedding dimension.\r\n - value: :math:`(S, E)` or :math:`(S, N, E)` where S is the source sequence length, N is the batch size, E is\r\n the embedding dimension.\r\n - key_padding_mask: :math:`(S)` or :math:`(N, S)` where N is the batch size, S is the source sequence length.\r\n If a FloatTensor is provided, it will be directly added to the value.\r\n If a BoolTensor is provided, the positions with the\r\n value of ``True`` will be ignored while the position with the value of ``False`` will be unchanged.\r\n - attn_mask: 2D mask :math:`(L, S)` where L is the target sequence length, S is the source sequence length.\r\n 3D mask :math:`(N*num_heads, L, S)` where N is the batch size, L is the target sequence length,\r\n S is the source sequence length. attn_mask ensures that position i is allowed to attend the unmasked\r\n positions. If a BoolTensor is provided, positions with ``True``\r\n are not allowed to attend while ``False`` values will be unchanged. If a FloatTensor\r\n is provided, it will be added to the attention weight.\r\n - static_k: :math:`(N*num_heads, S, E/num_heads)`, where S is the source sequence length,\r\n N is the batch size, E is the embedding dimension. E/num_heads is the head dimension.\r\n - static_v: :math:`(N*num_heads, S, E/num_heads)`, where S is the source sequence length,\r\n N is the batch size, E is the embedding dimension. E/num_heads is the head dimension.\r\n\r\n Outputs:\r\n - attn_output: :math:`(L, E)` or :math:`(L, N, E)` where L is the target sequence length, N is the batch size,\r\n E is the embedding dimension.\r\n - attn_output_weights: Only returned when ``need_weights=True``. If ``average_attn_weights=True``, returns\r\n attention weights averaged across heads of shape :math:`(L, S)` when input is unbatched or\r\n :math:`(N, L, S)`, where :math:`N` is the batch size, :math:`L` is the target sequence length, and\r\n :math:`S` is the source sequence length. If ``average_attn_weights=False``, returns attention weights per\r\n head of shape :math:`(num_heads, L, S)` when input is unbatched or :math:`(N, num_heads, L, S)`.\r\n \"\"\"\r\n tens_ops = (\r\n query,\r\n key,\r\n value,\r\n in_proj_weight,\r\n in_proj_bias,\r\n bias_k,\r\n bias_v,\r\n out_proj_weight,\r\n out_proj_bias,\r\n )\r\n if has_torch_function(tens_ops):\r\n return handle_torch_function(\r\n multi_head_attention_forward,\r\n tens_ops,\r\n query,\r\n key,\r\n value,\r\n embed_dim_to_check,\r\n num_heads,\r\n in_proj_weight,\r\n in_proj_bias,\r\n bias_k,\r\n bias_v,\r\n add_zero_attn,\r\n dropout_p,\r\n out_proj_weight,\r\n out_proj_bias,\r\n training=training,\r\n key_padding_mask=key_padding_mask,\r\n need_weights=need_weights,\r\n attn_mask=attn_mask,\r\n is_causal=is_causal,\r\n use_separate_proj_weight=use_separate_proj_weight,\r\n q_proj_weight=q_proj_weight,\r\n k_proj_weight=k_proj_weight,\r\n v_proj_weight=v_proj_weight,\r\n static_k=static_k,\r\n static_v=static_v,\r\n average_attn_weights=average_attn_weights,\r\n cache=cache,\r\n )\r\n\r\n is_batched = _mha_shape_check(query, key, value, key_padding_mask, attn_mask, num_heads)\r\n\r\n # For unbatched input, we unsqueeze at the expected batch-dim to pretend that the input\r\n # is batched, run the computation and before returning squeeze the\r\n # batch dimension so that the output doesn't carry this temporary batch dimension.\r\n if not is_batched:\r\n # unsqueeze if the input is unbatched\r\n query = query.unsqueeze(1)\r\n key = key.unsqueeze(1)\r\n value = value.unsqueeze(1)\r\n if key_padding_mask is not None:\r\n key_padding_mask = key_padding_mask.unsqueeze(0)\r\n\r\n # set up shape vars\r\n tgt_len, bsz, embed_dim = query.shape\r\n src_len, _, _ = key.shape\r\n\r\n key_padding_mask = _canonical_mask(\r\n mask=key_padding_mask,\r\n mask_name=\"key_padding_mask\",\r\n other_type=_none_or_dtype(attn_mask),\r\n other_name=\"attn_mask\",\r\n target_type=query.dtype,\r\n )\r\n\r\n if is_causal and attn_mask is None:\r\n raise RuntimeError(\r\n \"Need attn_mask if specifying the is_causal hint. \"\r\n \"You may use the Transformer module method \"\r\n \"`generate_square_subsequent_mask` to create this mask.\"\r\n )\r\n\r\n if is_causal and key_padding_mask is None and not need_weights:\r\n # when we have a kpm or need weights, we need attn_mask\r\n # Otherwise, we use the is_causal hint go as is_causal\r\n # indicator to SDPA.\r\n attn_mask = None\r\n else:\r\n attn_mask = _canonical_mask(\r\n mask=attn_mask,\r\n mask_name=\"attn_mask\",\r\n other_type=None,\r\n other_name=\"\",\r\n target_type=query.dtype,\r\n check_other=False,\r\n )\r\n\r\n if key_padding_mask is not None:\r\n # We have the attn_mask, and use that to merge kpm into it.\r\n # Turn off use of is_causal hint, as the merged mask is no\r\n # longer causal.\r\n is_causal = False\r\n\r\n assert embed_dim == embed_dim_to_check, (\r\n f\"was expecting embedding dimension of {embed_dim_to_check}, but got {embed_dim}\"\r\n )\r\n if isinstance(embed_dim, torch.Tensor):\r\n # embed_dim can be a tensor when JIT tracing\r\n head_dim = embed_dim.div(num_heads, rounding_mode=\"trunc\")\r\n else:\r\n head_dim = embed_dim // num_heads\r\n assert head_dim * num_heads == embed_dim, f\"embed_dim {embed_dim} not divisible by num_heads {num_heads}\"\r\n if use_separate_proj_weight:\r\n # allow MHA to have different embedding dimensions when separate projection weights are used\r\n assert key.shape[:2] == value.shape[:2], (\r\n f\"key's sequence and batch dims {key.shape[:2]} do not match value's {value.shape[:2]}\"\r\n )\r\n else:\r\n assert key.shape == value.shape, f\"key shape {key.shape} does not match value shape {value.shape}\"\r\n\r\n #\r\n # compute in-projection\r\n #\r\n if not use_separate_proj_weight:\r\n assert in_proj_weight is not None, \"use_separate_proj_weight is False but in_proj_weight is None\"\r\n q, k, v = _in_projection_packed(query, key, value, in_proj_weight, in_proj_bias)\r\n else:\r\n assert q_proj_weight is not None, \"use_separate_proj_weight is True but q_proj_weight is None\"\r\n assert k_proj_weight is not None, \"use_separate_proj_weight is True but k_proj_weight is None\"\r\n assert v_proj_weight is not None, \"use_separate_proj_weight is True but v_proj_weight is None\"\r\n if in_proj_bias is None:\r\n b_q = b_k = b_v = None\r\n else:\r\n b_q, b_k, b_v = in_proj_bias.chunk(3)\r\n q, k, v = _in_projection(\r\n query,\r\n key,\r\n value,\r\n q_proj_weight,\r\n k_proj_weight,\r\n v_proj_weight,\r\n b_q,\r\n b_k,\r\n b_v,\r\n )\r\n if cache != None:\r\n if cache[\"first_infer\"] == 1:\r\n cache[\"k\"][cache[\"stage\"]] = k\r\n # print(0,cache[\"k\"].shape)\r\n cache[\"v\"][cache[\"stage\"]] = v\r\n else: ###12个layer每个都要留自己的cache_kv\r\n # print(1,cache[\"k\"].shape)\r\n cache[\"k\"][cache[\"stage\"]] = torch.cat(\r\n [cache[\"k\"][cache[\"stage\"]], k], 0\r\n ) ##本来时序是1，但是proj的时候可能transpose了所以时序到0维了\r\n cache[\"v\"][cache[\"stage\"]] = torch.cat([cache[\"v\"][cache[\"stage\"]], v], 0)\r\n # print(2, cache[\"k\"].shape)\r\n src_len = cache[\"k\"][cache[\"stage\"]].shape[0]\r\n k = cache[\"k\"][cache[\"stage\"]]\r\n v = cache[\"v\"][cache[\"stage\"]]\r\n # if attn_mask is not None:\r\n # attn_mask=attn_mask[-1:,]\r\n # print(attn_mask.shape,attn_mask)\r\n cache[\"stage\"] = (cache[\"stage\"] + 1) % cache[\"all_stage\"]\r\n # print(2333,cache)\r\n # prep attention mask\r\n\r\n attn_mask = _canonical_mask(\r\n mask=attn_mask,\r\n mask_name=\"attn_mask\",\r\n other_type=None,\r\n other_name=\"\",\r\n target_type=q.dtype,\r\n check_other=False,\r\n )\r\n\r\n if attn_mask is not None:\r\n # ensure attn_mask's dim is 3\r\n if attn_mask.dim() == 2:\r\n correct_2d_size = (tgt_len, src_len)\r\n if attn_mask.shape != correct_2d_size:\r\n raise RuntimeError(\r\n f\"The shape of the 2D attn_mask is {attn_mask.shape}, but should be {correct_2d_size}.\"\r\n )\r\n attn_mask = attn_mask.unsqueeze(0)\r\n elif attn_mask.dim() == 3:\r\n correct_3d_size = (bsz * num_heads, tgt_len, src_len)\r\n if attn_mask.shape != correct_3d_size:\r\n raise RuntimeError(\r\n f\"The shape of the 3D attn_mask is {attn_mask.shape}, but should be {correct_3d_size}.\"\r\n )\r\n else:\r\n raise RuntimeError(f\"attn_mask's dimension {attn_mask.dim()} is not supported\")\r\n\r\n # add bias along batch dimension (currently second)\r\n if bias_k is not None and bias_v is not None:\r\n assert static_k is None, \"bias cannot be added to static key.\"\r\n assert static_v is None, \"bias cannot be added to static value.\"\r\n k = torch.cat([k, bias_k.repeat(1, bsz, 1)])\r\n v = torch.cat([v, bias_v.repeat(1, bsz, 1)])\r\n if attn_mask is not None:\r\n attn_mask = pad(attn_mask, (0, 1))\r\n if key_padding_mask is not None:\r\n key_padding_mask = pad(key_padding_mask, (0, 1))\r\n else:\r\n assert bias_k is None\r\n assert bias_v is None\r\n\r\n #\r\n # reshape q, k, v for multihead attention and make em batch first\r\n #\r\n q = q.view(tgt_len, bsz * num_heads, head_dim).transpose(0, 1)\r\n if static_k is None:\r\n k = k.view(k.shape[0], bsz * num_heads, head_dim).transpose(0, 1)\r\n else:\r\n # TODO finish disentangling control flow so we don't do in-projections when statics are passed\r\n assert static_k.size(0) == bsz * num_heads, (\r\n f\"expecting static_k.size(0) of {bsz * num_heads}, but got {static_k.size(0)}\"\r\n )\r\n assert static_k.size(2) == head_dim, f\"expecting static_k.size(2) of {head_dim}, but got {static_k.size(2)}\"\r\n k = static_k\r\n if static_v is None:\r\n v = v.view(v.shape[0], bsz * num_heads, head_dim).transpose(0, 1)\r\n else:\r\n # TODO finish disentangling control flow so we don't do in-projections when statics are passed\r\n assert static_v.size(0) == bsz * num_heads, (\r\n f\"expecting static_v.size(0) of {bsz * num_heads}, but got {static_v.size(0)}\"\r\n )\r\n assert static_v.size(2) == head_dim, f\"expecting static_v.size(2) of {head_dim}, but got {static_v.size(2)}\"\r\n v = static_v\r\n\r\n # add zero attention along batch dimension (now first)\r\n if add_zero_attn:\r\n zero_attn_shape = (bsz * num_heads, 1, head_dim)\r\n k = torch.cat([k, torch.zeros(zero_attn_shape, dtype=k.dtype, device=k.device)], dim=1)\r\n v = torch.cat([v, torch.zeros(zero_attn_shape, dtype=v.dtype, device=v.device)], dim=1)\r\n if attn_mask is not None:\r\n attn_mask = pad(attn_mask, (0, 1))\r\n if key_padding_mask is not None:\r\n key_padding_mask = pad(key_padding_mask, (0, 1))\r\n\r\n # update source sequence length after adjustments\r\n src_len = k.size(1)\r\n\r\n # merge key padding and attention masks\r\n if key_padding_mask is not None:\r\n assert key_padding_mask.shape == (\r\n bsz,\r\n src_len,\r\n ), f\"expecting key_padding_mask shape of {(bsz, src_len)}, but got {key_padding_mask.shape}\"\r\n key_padding_mask = (\r\n key_padding_mask.view(bsz, 1, 1, src_len).expand(-1, num_heads, -1, -1).reshape(bsz * num_heads, 1, src_len)\r\n )\r\n if attn_mask is None:\r\n attn_mask = key_padding_mask\r\n else:\r\n attn_mask = attn_mask + key_padding_mask\r\n\r\n # adjust dropout probability\r\n if not training:\r\n dropout_p = 0.0\r\n\r\n #\r\n # (deep breath) calculate attention and out projection\r\n #\r\n\r\n if need_weights:\r\n B, Nt, E = q.shape\r\n q_scaled = q / math.sqrt(E)\r\n\r\n assert not (is_causal and attn_mask is None), \"FIXME: is_causal not implemented for need_weights\"\r\n\r\n if attn_mask is not None:\r\n attn_output_weights = torch.baddbmm(attn_mask, q_scaled, k.transpose(-2, -1))\r\n else:\r\n attn_output_weights = torch.bmm(q_scaled, k.transpose(-2, -1))\r\n attn_output_weights = softmax(attn_output_weights, dim=-1)\r\n if dropout_p > 0.0:\r\n attn_output_weights = dropout(attn_output_weights, p=dropout_p)\r\n\r\n attn_output = torch.bmm(attn_output_weights, v)\r\n\r\n attn_output = attn_output.transpose(0, 1).contiguous().view(tgt_len * bsz, embed_dim)\r\n attn_output = linear(attn_output, out_proj_weight, out_proj_bias)\r\n attn_output = attn_output.view(tgt_len, bsz, attn_output.size(1))\r\n\r\n # optionally average attention weights over heads\r\n attn_output_weights = attn_output_weights.view(bsz, num_heads, tgt_len, src_len)\r\n if average_attn_weights:\r\n attn_output_weights = attn_output_weights.mean(dim=1)\r\n\r\n if not is_batched:\r\n # squeeze the output if input was unbatched\r\n attn_output = attn_output.squeeze(1)\r\n attn_output_weights = attn_output_weights.squeeze(0)\r\n return attn_output, attn_output_weights\r\n else:\r\n # attn_mask can be either (L,S) or (N*num_heads, L, S)\r\n # if attn_mask's shape is (1, L, S) we need to unsqueeze to (1, 1, L, S)\r\n # in order to match the input for SDPA of (N, num_heads, L, S)\r\n if attn_mask is not None:\r\n if attn_mask.size(0) == 1 and attn_mask.dim() == 3:\r\n attn_mask = attn_mask.unsqueeze(0)\r\n else:\r\n attn_mask = attn_mask.view(bsz, num_heads, -1, src_len)\r\n\r\n q = q.view(bsz, num_heads, tgt_len, head_dim)\r\n k = k.view(bsz, num_heads, src_len, head_dim)\r\n v = v.view(bsz, num_heads, src_len, head_dim)\r\n\r\n # with torch.backends.cuda.sdp_kernel(enable_flash=True, enable_math=True, enable_mem_efficient=True):\r\n attn_output = scaled_dot_product_attention(q, k, v, attn_mask, dropout_p, is_causal)\r\n\r\n attn_output = attn_output.permute(2, 0, 1, 3).contiguous().view(bsz * tgt_len, embed_dim)\r\n\r\n attn_output = linear(attn_output, out_proj_weight, out_proj_bias)\r\n attn_output = attn_output.view(tgt_len, bsz, attn_output.size(1))\r\n if not is_batched:\r\n # squeeze the output if input was unbatched\r\n attn_output = attn_output.squeeze(1)\r\n return attn_output, None\r\n"
},
{
"path": "GPT_SoVITS/AR/modules/patched_mha_with_cache_onnx.py",
"content": "from torch.nn.functional import *\r\nfrom torch.nn.functional import (\r\n _canonical_mask,\r\n)\r\n\r\n\r\ndef multi_head_attention_forward_patched(\r\n query,\r\n key,\r\n value,\r\n embed_dim_to_check: int,\r\n num_heads: int,\r\n in_proj_weight,\r\n in_proj_bias: Optional[Tensor],\r\n bias_k: Optional[Tensor],\r\n bias_v: Optional[Tensor],\r\n add_zero_attn: bool,\r\n dropout_p: float,\r\n out_proj_weight: Tensor,\r\n out_proj_bias: Optional[Tensor],\r\n training: bool = True,\r\n key_padding_mask: Optional[Tensor] = None,\r\n need_weights: bool = True,\r\n attn_mask: Optional[Tensor] = None,\r\n use_separate_proj_weight: bool = False,\r\n q_proj_weight: Optional[Tensor] = None,\r\n k_proj_weight: Optional[Tensor] = None,\r\n v_proj_weight: Optional[Tensor] = None,\r\n static_k: Optional[Tensor] = None,\r\n static_v: Optional[Tensor] = None,\r\n average_attn_weights: bool = True,\r\n is_causal: bool = False,\r\n cache=None,\r\n) -> Tuple[Tensor, Optional[Tensor]]:\r\n # set up shape vars\r\n _, _, embed_dim = query.shape\r\n attn_mask = _canonical_mask(\r\n mask=attn_mask,\r\n mask_name=\"attn_mask\",\r\n other_type=None,\r\n other_name=\"\",\r\n target_type=query.dtype,\r\n check_other=False,\r\n )\r\n head_dim = embed_dim // num_heads\r\n\r\n proj_qkv = linear(query, in_proj_weight, in_proj_bias)\r\n proj_qkv = proj_qkv.unflatten(-1, (3, query.size(-1))).unsqueeze(0).transpose(0, -2).squeeze(-2).contiguous()\r\n q, k, v = proj_qkv[0], proj_qkv[1], proj_qkv[2]\r\n\r\n if cache[\"first_infer\"] == 1:\r\n cache[\"k\"][cache[\"stage\"]] = k\r\n cache[\"v\"][cache[\"stage\"]] = v\r\n else:\r\n cache[\"k\"][cache[\"stage\"]] = torch.cat([cache[\"k\"][cache[\"stage\"]][:-1], k], 0)\r\n cache[\"v\"][cache[\"stage\"]] = torch.cat([cache[\"v\"][cache[\"stage\"]][:-1], v], 0)\r\n k = cache[\"k\"][cache[\"stage\"]]\r\n v = cache[\"v\"][cache[\"stage\"]]\r\n cache[\"stage\"] = (cache[\"stage\"] + 1) % cache[\"all_stage\"]\r\n\r\n attn_mask = _canonical_mask(\r\n mask=attn_mask,\r\n mask_name=\"attn_mask\",\r\n other_type=None,\r\n other_name=\"\",\r\n target_type=q.dtype,\r\n check_other=False,\r\n )\r\n attn_mask = attn_mask.unsqueeze(0)\r\n\r\n q = q.view(-1, num_heads, head_dim).transpose(0, 1)\r\n k = k.view(-1, num_heads, head_dim).transpose(0, 1)\r\n v = v.view(-1, num_heads, head_dim).transpose(0, 1)\r\n\r\n dropout_p = 0.0\r\n attn_mask = attn_mask.unsqueeze(0)\r\n q = q.view(num_heads, -1, head_dim).unsqueeze(0)\r\n k = k.view(num_heads, -1, head_dim).unsqueeze(0)\r\n v = v.view(num_heads, -1, head_dim).unsqueeze(0)\r\n attn_output = scaled_dot_product_attention(q, k, v, attn_mask, dropout_p, is_causal)\r\n attn_output = attn_output.permute(2, 0, 1, 3).contiguous().view(-1, embed_dim)\r\n attn_output = linear(attn_output, out_proj_weight, out_proj_bias)\r\n attn_output = attn_output.view(-1, 1, attn_output.size(1))\r\n\r\n return attn_output\r\n"
},
{
"path": "GPT_SoVITS/AR/modules/scaling.py",
"content": "# Copyright 2022 Xiaomi Corp. (authors: Daniel Povey)\n#\n# See ../../../../LICENSE for clarification regarding multiple authors\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\nimport random\nfrom typing import Optional\nfrom typing import Tuple\n\nimport torch\nimport torch.nn as nn\nfrom torch import Tensor\n\n\nclass DoubleSwishFunction(torch.autograd.Function):\n \"\"\"\n double_swish(x) = x * torch.sigmoid(x-1)\n This is a definition, originally motivated by its close numerical\n similarity to swish(swish(x)), where swish(x) = x * sigmoid(x).\n\n Memory-efficient derivative computation:\n double_swish(x) = x * s, where s(x) = torch.sigmoid(x-1)\n double_swish'(x) = d/dx double_swish(x) = x * s'(x) + x' * s(x) = x * s'(x) + s(x).\n Now, s'(x) = s(x) * (1-s(x)).\n double_swish'(x) = x * s'(x) + s(x).\n = x * s(x) * (1-s(x)) + s(x).\n = double_swish(x) * (1-s(x)) + s(x)\n ... so we just need to remember s(x) but not x itself.\n \"\"\"\n\n @staticmethod\n def forward(ctx, x: Tensor) -> Tensor:\n requires_grad = x.requires_grad\n x_dtype = x.dtype\n if x.dtype == torch.float16:\n x = x.to(torch.float32)\n\n s = torch.sigmoid(x - 1.0)\n y = x * s\n\n if requires_grad:\n deriv = y * (1 - s) + s\n # notes on derivative of x * sigmoid(x - 1):\n # https://www.wolframalpha.com/input?i=d%2Fdx+%28x+*+sigmoid%28x-1%29%29\n # min \\simeq -0.043638. Take floor as -0.043637 so it's a lower bund\n # max \\simeq 1.1990. Take ceil to be 1.2 so it's an upper bound.\n # the combination of \"+ torch.rand_like(deriv)\" and casting to torch.uint8 (which\n # floors), should be expectation-preserving.\n floor = -0.043637\n ceil = 1.2\n d_scaled = (deriv - floor) * (255.0 / (ceil - floor)) + torch.rand_like(deriv)\n if __name__ == \"__main__\":\n # for self-testing only.\n assert d_scaled.min() >= 0.0\n assert d_scaled.max() < 256.0\n d_int = d_scaled.to(torch.uint8)\n ctx.save_for_backward(d_int)\n if x.dtype == torch.float16 or torch.is_autocast_enabled():\n y = y.to(torch.float16)\n return y\n\n @staticmethod\n def backward(ctx, y_grad: Tensor) -> Tensor:\n (d,) = ctx.saved_tensors\n # the same constants as used in forward pass.\n floor = -0.043637\n ceil = 1.2\n d = d * ((ceil - floor) / 255.0) + floor\n return y_grad * d\n\n\nclass DoubleSwish(torch.nn.Module):\n def forward(self, x: Tensor) -> Tensor:\n \"\"\"Return double-swish activation function which is an approximation to Swish(Swish(x)),\n that we approximate closely with x * sigmoid(x-1).\n \"\"\"\n if torch.jit.is_scripting() or torch.jit.is_tracing():\n return x * torch.sigmoid(x - 1.0)\n return DoubleSwishFunction.apply(x)\n\n\nclass ActivationBalancerFunction(torch.autograd.Function):\n @staticmethod\n def forward(\n ctx,\n x: Tensor,\n scale_factor: Tensor,\n sign_factor: Optional[Tensor],\n channel_dim: int,\n ) -> Tensor:\n if channel_dim < 0:\n channel_dim += x.ndim\n ctx.channel_dim = channel_dim\n xgt0 = x > 0\n if sign_factor is None:\n ctx.save_for_backward(xgt0, scale_factor)\n else:\n ctx.save_for_backward(xgt0, scale_factor, sign_factor)\n return x\n\n @staticmethod\n def backward(ctx, x_grad: Tensor) -> Tuple[Tensor, None, None, None]:\n if len(ctx.saved_tensors) == 3:\n xgt0, scale_factor, sign_factor = ctx.saved_tensors\n for _ in range(ctx.channel_dim, x_grad.ndim - 1):\n scale_factor = scale_factor.unsqueeze(-1)\n sign_factor = sign_factor.unsqueeze(-1)\n factor = sign_factor + scale_factor * (xgt0.to(x_grad.dtype) - 0.5)\n else:\n xgt0, scale_factor = ctx.saved_tensors\n for _ in range(ctx.channel_dim, x_grad.ndim - 1):\n scale_factor = scale_factor.unsqueeze(-1)\n factor = scale_factor * (xgt0.to(x_grad.dtype) - 0.5)\n neg_delta_grad = x_grad.abs() * factor\n return (\n x_grad - neg_delta_grad,\n None,\n None,\n None,\n )\n\n\ndef _compute_scale_factor(\n x: Tensor,\n channel_dim: int,\n min_abs: float,\n max_abs: float,\n gain_factor: float,\n max_factor: float,\n) -> Tensor:\n if channel_dim < 0:\n channel_dim += x.ndim\n sum_dims = [d for d in range(x.ndim) if d != channel_dim]\n x_abs_mean = torch.mean(x.abs(), dim=sum_dims).to(torch.float32)\n\n if min_abs == 0.0:\n below_threshold = 0.0\n else:\n # below_threshold is 0 if x_abs_mean > min_abs, can be at most max_factor if\n # x_abs)_mean , min_abs.\n below_threshold = ((min_abs - x_abs_mean) * (gain_factor / min_abs)).clamp(min=0, max=max_factor)\n\n above_threshold = ((x_abs_mean - max_abs) * (gain_factor / max_abs)).clamp(min=0, max=max_factor)\n\n return below_threshold - above_threshold\n\n\ndef _compute_sign_factor(\n x: Tensor,\n channel_dim: int,\n min_positive: float,\n max_positive: float,\n gain_factor: float,\n max_factor: float,\n) -> Tensor:\n if channel_dim < 0:\n channel_dim += x.ndim\n sum_dims = [d for d in range(x.ndim) if d != channel_dim]\n proportion_positive = torch.mean((x > 0).to(torch.float32), dim=sum_dims)\n if min_positive == 0.0:\n factor1 = 0.0\n else:\n # 0 if proportion_positive >= min_positive, else can be\n # as large as max_factor.\n factor1 = ((min_positive - proportion_positive) * (gain_factor / min_positive)).clamp_(min=0, max=max_factor)\n\n if max_positive == 1.0:\n factor2 = 0.0\n else:\n # 0 if self.proportion_positive <= max_positive, else can be\n # as large as -max_factor.\n factor2 = ((proportion_positive - max_positive) * (gain_factor / (1.0 - max_positive))).clamp_(\n min=0, max=max_factor\n )\n sign_factor = factor1 - factor2\n # require min_positive != 0 or max_positive != 1:\n assert not isinstance(sign_factor, float)\n return sign_factor\n\n\nclass ActivationBalancer(torch.nn.Module):\n \"\"\"\n Modifies the backpropped derivatives of a function to try to encourage, for\n each channel, that it is positive at least a proportion `threshold` of the\n time. It does this by multiplying negative derivative values by up to\n (1+max_factor), and positive derivative values by up to (1-max_factor),\n interpolated from 1 at the threshold to those extremal values when none\n of the inputs are positive.\n\n Args:\n num_channels: the number of channels\n channel_dim: the dimension/axis corresponding to the channel, e.g.\n -1, 0, 1, 2; will be interpreted as an offset from x.ndim if negative.\n min_positive: the minimum, per channel, of the proportion of the time\n that (x > 0), below which we start to modify the derivatives.\n max_positive: the maximum, per channel, of the proportion of the time\n that (x > 0), above which we start to modify the derivatives.\n max_factor: the maximum factor by which we modify the derivatives for\n either the sign constraint or the magnitude constraint;\n e.g. with max_factor=0.02, the the derivatives would be multiplied by\n values in the range [0.98..1.02].\n sign_gain_factor: determines the 'gain' with which we increase the\n change in gradient once the constraints on min_positive and max_positive\n are violated.\n scale_gain_factor: determines the 'gain' with which we increase the\n change in gradient once the constraints on min_abs and max_abs\n are violated.\n min_abs: the minimum average-absolute-value difference from the mean\n value per channel, which we allow, before we start to modify\n the derivatives to prevent this.\n max_abs: the maximum average-absolute-value difference from the mean\n value per channel, which we allow, before we start to modify\n the derivatives to prevent this.\n min_prob: determines the minimum probability with which we modify the\n gradients for the {min,max}_positive and {min,max}_abs constraints,\n on each forward(). This is done randomly to prevent all layers\n from doing it at the same time. Early in training we may use\n higher probabilities than this; it will decay to this value.\n \"\"\"\n\n def __init__(\n self,\n num_channels: int,\n channel_dim: int,\n min_positive: float = 0.05,\n max_positive: float = 0.95,\n max_factor: float = 0.04,\n sign_gain_factor: float = 0.01,\n scale_gain_factor: float = 0.02,\n min_abs: float = 0.2,\n max_abs: float = 100.0,\n min_prob: float = 0.1,\n ):\n super(ActivationBalancer, self).__init__()\n self.num_channels = num_channels\n self.channel_dim = channel_dim\n self.min_positive = min_positive\n self.max_positive = max_positive\n self.max_factor = max_factor\n self.min_abs = min_abs\n self.max_abs = max_abs\n self.min_prob = min_prob\n self.sign_gain_factor = sign_gain_factor\n self.scale_gain_factor = scale_gain_factor\n\n # count measures how many times the forward() function has been called.\n # We occasionally sync this to a tensor called `count`, that exists to\n # make sure it is synced to disk when we load and save the model.\n self.cpu_count = 0\n self.register_buffer(\"count\", torch.tensor(0, dtype=torch.int64))\n\n def forward(self, x: Tensor) -> Tensor:\n if torch.jit.is_scripting() or not x.requires_grad or torch.jit.is_tracing():\n return _no_op(x)\n\n count = self.cpu_count\n self.cpu_count += 1\n\n if random.random() < 0.01:\n # Occasionally sync self.cpu_count with self.count.\n # count affects the decay of 'prob'. don't do this on every iter,\n # because syncing with the GPU is slow.\n self.cpu_count = max(self.cpu_count, self.count.item())\n self.count.fill_(self.cpu_count)\n\n # the prob of doing some work exponentially decreases from 0.5 till it hits\n # a floor at min_prob (==0.1, by default)\n prob = max(self.min_prob, 0.5 ** (1 + (count / 4000.0)))\n\n if random.random() < prob:\n sign_gain_factor = 0.5\n if self.min_positive != 0.0 or self.max_positive != 1.0:\n sign_factor = _compute_sign_factor(\n x,\n self.channel_dim,\n self.min_positive,\n self.max_positive,\n gain_factor=self.sign_gain_factor / prob,\n max_factor=self.max_factor,\n )\n else:\n sign_factor = None\n\n scale_factor = _compute_scale_factor(\n x.detach(),\n self.channel_dim,\n min_abs=self.min_abs,\n max_abs=self.max_abs,\n gain_factor=self.scale_gain_factor / prob,\n max_factor=self.max_factor,\n )\n return ActivationBalancerFunction.apply(\n x,\n scale_factor,\n sign_factor,\n self.channel_dim,\n )\n else:\n return _no_op(x)\n\n\ndef BalancedDoubleSwish(d_model, channel_dim=-1, max_abs=10.0, min_prob=0.25) -> nn.Sequential:\n \"\"\"\n ActivationBalancer -> DoubleSwish\n \"\"\"\n balancer = ActivationBalancer(d_model, channel_dim=channel_dim, max_abs=max_abs, min_prob=min_prob)\n return nn.Sequential(\n balancer,\n DoubleSwish(),\n )\n"
},
{
"path": "GPT_SoVITS/AR/modules/transformer.py",
"content": "# modified from https://github.com/lifeiteng/vall-e/blob/main/valle/modules/transformer.py\nimport copy\nimport numbers\nfrom functools import partial\nfrom typing import Any\nfrom typing import Callable\nfrom typing import List\nfrom typing import Optional\nfrom typing import Tuple\nfrom typing import Union\n\nimport torch\nfrom AR.modules.activation import MultiheadAttention\nfrom AR.modules.scaling import BalancedDoubleSwish\nfrom torch import nn\nfrom torch import Tensor\nfrom torch.nn import functional as F\n\n_shape_t = Union[int, List[int], torch.Size]\n\n\nclass LayerNorm(nn.Module):\n __constants__ = [\"normalized_shape\", \"eps\", \"elementwise_affine\"]\n normalized_shape: Tuple[int, ...]\n eps: float\n elementwise_affine: bool\n\n def __init__(\n self,\n normalized_shape: _shape_t,\n eps: float = 1e-5,\n elementwise_affine: bool = True,\n device=None,\n dtype=None,\n ) -> None:\n factory_kwargs = {\"device\": device, \"dtype\": dtype}\n super(LayerNorm, self).__init__()\n if isinstance(normalized_shape, numbers.Integral):\n # mypy error: incompatible types in assignment\n normalized_shape = (normalized_shape,) # type: ignore[assignment]\n self.normalized_shape = tuple(normalized_shape) # type: ignore[arg-type]\n self.eps = eps\n self.elementwise_affine = elementwise_affine\n if self.elementwise_affine:\n self.weight = nn.Parameter(torch.empty(self.normalized_shape, **factory_kwargs))\n self.bias = nn.Parameter(torch.empty(self.normalized_shape, **factory_kwargs))\n else:\n self.register_parameter(\"weight\", None)\n self.register_parameter(\"bias\", None)\n\n self.reset_parameters()\n\n def reset_parameters(self) -> None:\n if self.elementwise_affine:\n nn.init.ones_(self.weight)\n nn.init.zeros_(self.bias)\n\n def forward(self, input: Tensor, embedding: Any = None) -> Tensor:\n if isinstance(input, tuple):\n input, embedding = input\n return (\n F.layer_norm(\n input,\n self.normalized_shape,\n self.weight,\n self.bias,\n self.eps,\n ),\n embedding,\n )\n\n assert embedding is None\n return F.layer_norm(input, self.normalized_shape, self.weight, self.bias, self.eps)\n\n def extra_repr(self) -> str:\n return \"{normalized_shape}, eps={eps}, elementwise_affine={elementwise_affine}\".format(**self.__dict__)\n\n\nclass IdentityNorm(nn.Module):\n def __init__(\n self,\n d_model: int,\n eps: float = 1e-5,\n device=None,\n dtype=None,\n ) -> None:\n super(IdentityNorm, self).__init__()\n\n def forward(self, input: Tensor, embedding: Any = None) -> Tensor:\n if isinstance(input, tuple):\n return input\n\n assert embedding is None\n return input\n\n\nclass TransformerEncoder(nn.Module):\n r\"\"\"TransformerEncoder is a stack of N encoder layers. Users can build the\n BERT(https://arxiv.org/abs/1810.04805) model with corresponding parameters.\n\n Args:\n encoder_layer: an instance of the TransformerEncoderLayer() class (required).\n num_layers: the number of sub-encoder-layers in the encoder (required).\n norm: the layer normalization component (optional).\n enable_nested_tensor: if True, input will automatically convert to nested tensor\n (and convert back on output). This will improve the overall performance of\n TransformerEncoder when padding rate is high. Default: ``True`` (enabled).\n\n Examples::\n >>> encoder_layer = TransformerEncoderLayer(d_model=512, nhead=8)\n >>> transformer_encoder = TransformerEncoder(encoder_layer, num_layers=6)\n >>> src = torch.rand(10, 32, 512)\n >>> out = transformer_encoder(src)\n \"\"\"\n\n __constants__ = [\"norm\"]\n\n def __init__(self, encoder_layer, num_layers, norm=None):\n super(TransformerEncoder, self).__init__()\n self.layers = _get_clones(encoder_layer, num_layers)\n self.num_layers = num_layers\n self.norm = norm\n\n def forward(\n self,\n src: Tensor,\n mask: Optional[Tensor] = None,\n src_key_padding_mask: Optional[Tensor] = None,\n return_layer_states: bool = False,\n cache=None,\n ) -> Tensor:\n r\"\"\"Pass the input through the encoder layers in turn.\n\n Args:\n src: the sequence to the encoder (required).\n mask: the mask for the src sequence (optional).\n src_key_padding_mask: the mask for the src keys per batch (optional).\n return_layer_states: return layers' state (optional).\n\n Shape:\n see the docs in Transformer class.\n \"\"\"\n if return_layer_states:\n layer_states = [] # layers' output\n output = src\n for mod in self.layers:\n output = mod(\n output,\n src_mask=mask,\n src_key_padding_mask=src_key_padding_mask,\n cache=cache,\n )\n layer_states.append(output[0])\n\n if self.norm is not None:\n output = self.norm(output)\n\n return layer_states, output\n\n output = src\n for mod in self.layers:\n output = mod(\n output,\n src_mask=mask,\n src_key_padding_mask=src_key_padding_mask,\n cache=cache,\n )\n\n if self.norm is not None:\n output = self.norm(output)\n\n return output\n\n\nclass TransformerEncoderLayer(nn.Module):\n __constants__ = [\"batch_first\", \"norm_first\"]\n\n def __init__(\n self,\n d_model: int,\n nhead: int,\n dim_feedforward: int = 2048,\n dropout: float = 0.1,\n activation: Union[str, Callable[[Tensor], Tensor]] = F.relu,\n batch_first: bool = False,\n norm_first: bool = False,\n device=None,\n dtype=None,\n linear1_self_attention_cls: nn.Module = nn.Linear,\n linear2_self_attention_cls: nn.Module = nn.Linear,\n linear1_feedforward_cls: nn.Module = nn.Linear,\n linear2_feedforward_cls: nn.Module = nn.Linear,\n layer_norm_cls: nn.Module = LayerNorm,\n layer_norm_eps: float = 1e-5,\n adaptive_layer_norm=False,\n ) -> None:\n factory_kwargs = {\"device\": device, \"dtype\": dtype}\n super(TransformerEncoderLayer, self).__init__()\n # print(233333333333,d_model,nhead)\n # import os\n # os._exit(2333333)\n self.self_attn = MultiheadAttention(\n d_model, # 512 16\n nhead,\n dropout=dropout,\n batch_first=batch_first,\n linear1_cls=linear1_self_attention_cls,\n linear2_cls=linear2_self_attention_cls,\n **factory_kwargs,\n )\n\n # Implementation of Feedforward model\n self.linear1 = linear1_feedforward_cls(d_model, dim_feedforward, **factory_kwargs)\n self.dropout = nn.Dropout(dropout)\n self.linear2 = linear2_feedforward_cls(dim_feedforward, d_model, **factory_kwargs)\n\n self.norm_first = norm_first\n self.dropout1 = nn.Dropout(dropout)\n self.dropout2 = nn.Dropout(dropout)\n\n # Legacy string support for activation function.\n if isinstance(activation, str):\n activation = _get_activation_fn(activation)\n elif isinstance(activation, partial):\n activation = activation(d_model)\n elif activation == BalancedDoubleSwish:\n activation = BalancedDoubleSwish(d_model)\n\n # # We can't test self.activation in forward() in TorchScript,\n # # so stash some information about it instead.\n # if activation is F.relu or isinstance(activation, torch.nn.ReLU):\n # self.activation_relu_or_gelu = 1\n # elif activation is F.gelu or isinstance(activation, torch.nn.GELU):\n # self.activation_relu_or_gelu = 2\n # else:\n # self.activation_relu_or_gelu = 0\n self.activation = activation\n\n norm1 = layer_norm_cls(d_model, eps=layer_norm_eps, **factory_kwargs)\n if layer_norm_cls == IdentityNorm:\n norm2 = BalancedBasicNorm(d_model, eps=layer_norm_eps, **factory_kwargs)\n else:\n norm2 = layer_norm_cls(d_model, eps=layer_norm_eps, **factory_kwargs)\n\n if adaptive_layer_norm:\n self.norm1 = AdaptiveLayerNorm(d_model, norm1)\n self.norm2 = AdaptiveLayerNorm(d_model, norm2)\n else:\n self.norm1 = norm1\n self.norm2 = norm2\n\n def __setstate__(self, state):\n super(TransformerEncoderLayer, self).__setstate__(state)\n if not hasattr(self, \"activation\"):\n self.activation = F.relu\n\n def forward(\n self,\n src: Tensor,\n src_mask: Optional[Tensor] = None,\n src_key_padding_mask: Optional[Tensor] = None,\n cache=None,\n ) -> Tensor:\n r\"\"\"Pass the input through the encoder layer.\n\n Args:\n src: the sequence to the encoder layer (required).\n src_mask: the mask for the src sequence (optional).\n src_key_padding_mask: the mask for the src keys per batch (optional).\n\n Shape:\n see the docs in Transformer class.\n \"\"\"\n x, stage_embedding = src, None\n is_src_tuple = False\n if isinstance(src, tuple):\n x, stage_embedding = src\n is_src_tuple = True\n\n if src_key_padding_mask is not None:\n _skpm_dtype = src_key_padding_mask.dtype\n if _skpm_dtype != torch.bool and not torch.is_floating_point(src_key_padding_mask):\n raise AssertionError(\"only bool and floating types of key_padding_mask are supported\")\n\n if self.norm_first:\n x = x + self._sa_block(\n self.norm1(x, stage_embedding),\n src_mask,\n src_key_padding_mask,\n cache=cache,\n )\n x = x + self._ff_block(self.norm2(x, stage_embedding))\n else:\n x = self.norm1(\n x + self._sa_block(x, src_mask, src_key_padding_mask, cache=cache),\n stage_embedding,\n )\n x = self.norm2(x + self._ff_block(x), stage_embedding)\n\n if is_src_tuple:\n return (x, stage_embedding)\n return x\n\n # self-attention block\n def _sa_block(\n self,\n x: Tensor,\n attn_mask: Optional[Tensor],\n key_padding_mask: Optional[Tensor],\n cache=None,\n ) -> Tensor:\n # print(x.shape,attn_mask.shape,key_padding_mask)\n # torch.Size([1, 188, 512]) torch.Size([188, 188]) None\n # import os\n # os._exit(23333)\n x = self.self_attn(\n x,\n x,\n x,\n attn_mask=attn_mask,\n key_padding_mask=key_padding_mask,\n need_weights=False,\n cache=cache,\n )[0]\n return self.dropout1(x)\n\n # feed forward block\n def _ff_block(self, x: Tensor) -> Tensor:\n x = self.linear2(self.dropout(self.activation(self.linear1(x))))\n return self.dropout2(x)\n\n\nclass AdaptiveLayerNorm(nn.Module):\n r\"\"\"Adaptive Layer Normalization\"\"\"\n\n def __init__(self, d_model, norm) -> None:\n super(AdaptiveLayerNorm, self).__init__()\n self.project_layer = nn.Linear(d_model, 2 * d_model)\n self.norm = norm\n self.d_model = d_model\n self.eps = self.norm.eps\n\n def forward(self, input: Tensor, embedding: Tensor = None) -> Tensor:\n if isinstance(input, tuple):\n input, embedding = input\n weight, bias = torch.split(\n self.project_layer(embedding),\n split_size_or_sections=self.d_model,\n dim=-1,\n )\n return (weight * self.norm(input) + bias, embedding)\n\n weight, bias = torch.split(\n self.project_layer(embedding),\n split_size_or_sections=self.d_model,\n dim=-1,\n )\n return weight * self.norm(input) + bias\n\n\ndef _get_clones(module, N):\n return nn.ModuleList([copy.deepcopy(module) for i in range(N)])\n"
},
{
"path": "GPT_SoVITS/AR/modules/transformer_onnx.py",
"content": "# modified from https://github.com/lifeiteng/vall-e/blob/main/valle/modules/transformer.py\nimport copy\nimport numbers\nfrom functools import partial\nfrom typing import Any\nfrom typing import Callable\nfrom typing import List\nfrom typing import Optional\nfrom typing import Tuple\nfrom typing import Union\n\nimport torch\nfrom AR.modules.activation_onnx import MultiheadAttention\nfrom AR.modules.scaling import BalancedDoubleSwish\nfrom torch import nn\nfrom torch import Tensor\nfrom torch.nn import functional as F\n\n_shape_t = Union[int, List[int], torch.Size]\n\n\nclass LayerNorm(nn.Module):\n __constants__ = [\"normalized_shape\", \"eps\", \"elementwise_affine\"]\n normalized_shape: Tuple[int, ...]\n eps: float\n elementwise_affine: bool\n\n def __init__(\n self,\n normalized_shape: _shape_t,\n eps: float = 1e-5,\n elementwise_affine: bool = True,\n device=None,\n dtype=None,\n ) -> None:\n factory_kwargs = {\"device\": device, \"dtype\": dtype}\n super(LayerNorm, self).__init__()\n if isinstance(normalized_shape, numbers.Integral):\n # mypy error: incompatible types in assignment\n normalized_shape = (normalized_shape,) # type: ignore[assignment]\n self.normalized_shape = tuple(normalized_shape) # type: ignore[arg-type]\n self.eps = eps\n self.elementwise_affine = elementwise_affine\n if self.elementwise_affine:\n self.weight = nn.Parameter(torch.empty(self.normalized_shape, **factory_kwargs))\n self.bias = nn.Parameter(torch.empty(self.normalized_shape, **factory_kwargs))\n else:\n self.register_parameter(\"weight\", None)\n self.register_parameter(\"bias\", None)\n\n self.reset_parameters()\n\n def reset_parameters(self) -> None:\n if self.elementwise_affine:\n nn.init.ones_(self.weight)\n nn.init.zeros_(self.bias)\n\n def forward(self, input: Tensor, embedding: Any = None) -> Tensor:\n if isinstance(input, tuple):\n input, embedding = input\n return (\n F.layer_norm(\n input,\n self.normalized_shape,\n self.weight,\n self.bias,\n self.eps,\n ),\n embedding,\n )\n\n assert embedding is None\n return F.layer_norm(input, self.normalized_shape, self.weight, self.bias, self.eps)\n\n def extra_repr(self) -> str:\n return \"{normalized_shape}, eps={eps}, elementwise_affine={elementwise_affine}\".format(**self.__dict__)\n\n\nclass IdentityNorm(nn.Module):\n def __init__(\n self,\n d_model: int,\n eps: float = 1e-5,\n device=None,\n dtype=None,\n ) -> None:\n super(IdentityNorm, self).__init__()\n\n def forward(self, input: Tensor, embedding: Any = None) -> Tensor:\n if isinstance(input, tuple):\n return input\n\n assert embedding is None\n return input\n\n\nclass TransformerEncoder(nn.Module):\n r\"\"\"TransformerEncoder is a stack of N encoder layers. Users can build the\n BERT(https://arxiv.org/abs/1810.04805) model with corresponding parameters.\n\n Args:\n encoder_layer: an instance of the TransformerEncoderLayer() class (required).\n num_layers: the number of sub-encoder-layers in the encoder (required).\n norm: the layer normalization component (optional).\n enable_nested_tensor: if True, input will automatically convert to nested tensor\n (and convert back on output). This will improve the overall performance of\n TransformerEncoder when padding rate is high. Default: ``True`` (enabled).\n\n Examples::\n >>> encoder_layer = TransformerEncoderLayer(d_model=512, nhead=8)\n >>> transformer_encoder = TransformerEncoder(encoder_layer, num_layers=6)\n >>> src = torch.rand(10, 32, 512)\n >>> out = transformer_encoder(src)\n \"\"\"\n\n __constants__ = [\"norm\"]\n\n def __init__(self, encoder_layer, num_layers, norm=None):\n super(TransformerEncoder, self).__init__()\n self.layers = _get_clones(encoder_layer, num_layers)\n self.num_layers = num_layers\n self.norm = norm\n\n def forward(\n self,\n src: Tensor,\n mask: Optional[Tensor] = None,\n src_key_padding_mask: Optional[Tensor] = None,\n return_layer_states: bool = False,\n cache=None,\n ) -> Tensor:\n output = src\n for mod in self.layers:\n output = mod(\n output,\n src_mask=mask,\n src_key_padding_mask=src_key_padding_mask,\n cache=cache,\n )\n\n if self.norm is not None:\n output = self.norm(output)\n\n return output\n\n\nclass TransformerEncoderLayer(nn.Module):\n __constants__ = [\"batch_first\", \"norm_first\"]\n\n def __init__(\n self,\n d_model: int,\n nhead: int,\n dim_feedforward: int = 2048,\n dropout: float = 0.1,\n activation: Union[str, Callable[[Tensor], Tensor]] = F.relu,\n batch_first: bool = False,\n norm_first: bool = False,\n device=None,\n dtype=None,\n linear1_self_attention_cls: nn.Module = nn.Linear,\n linear2_self_attention_cls: nn.Module = nn.Linear,\n linear1_feedforward_cls: nn.Module = nn.Linear,\n linear2_feedforward_cls: nn.Module = nn.Linear,\n layer_norm_cls: nn.Module = LayerNorm,\n layer_norm_eps: float = 1e-5,\n adaptive_layer_norm=False,\n ) -> None:\n factory_kwargs = {\"device\": device, \"dtype\": dtype}\n super(TransformerEncoderLayer, self).__init__()\n self.self_attn = MultiheadAttention(\n d_model, # 512 16\n nhead,\n dropout=dropout,\n batch_first=batch_first,\n linear1_cls=linear1_self_attention_cls,\n linear2_cls=linear2_self_attention_cls,\n **factory_kwargs,\n )\n self.linear1 = linear1_feedforward_cls(d_model, dim_feedforward, **factory_kwargs)\n self.dropout = nn.Dropout(dropout)\n self.linear2 = linear2_feedforward_cls(dim_feedforward, d_model, **factory_kwargs)\n self.norm_first = norm_first\n self.dropout1 = nn.Dropout(dropout)\n self.dropout2 = nn.Dropout(dropout)\n if isinstance(activation, str):\n activation = _get_activation_fn(activation)\n elif isinstance(activation, partial):\n activation = activation(d_model)\n elif activation == BalancedDoubleSwish:\n activation = BalancedDoubleSwish(d_model)\n self.activation = activation\n\n norm1 = layer_norm_cls(d_model, eps=layer_norm_eps, **factory_kwargs)\n if layer_norm_cls == IdentityNorm:\n norm2 = BalancedBasicNorm(d_model, eps=layer_norm_eps, **factory_kwargs)\n else:\n norm2 = layer_norm_cls(d_model, eps=layer_norm_eps, **factory_kwargs)\n\n if adaptive_layer_norm:\n self.norm1 = AdaptiveLayerNorm(d_model, norm1)\n self.norm2 = AdaptiveLayerNorm(d_model, norm2)\n else:\n self.norm1 = norm1\n self.norm2 = norm2\n\n def __setstate__(self, state):\n super(TransformerEncoderLayer, self).__setstate__(state)\n if not hasattr(self, \"activation\"):\n self.activation = F.relu\n\n def forward(\n self,\n src: Tensor,\n src_mask: Optional[Tensor] = None,\n src_key_padding_mask: Optional[Tensor] = None,\n cache=None,\n ) -> Tensor:\n x = src\n stage_embedding = None\n x = self.norm1(\n x + self._sa_block(x, src_mask, src_key_padding_mask, cache=cache),\n stage_embedding,\n )\n x = self.norm2(x + self._ff_block(x), stage_embedding)\n\n return x\n\n def _sa_block(\n self,\n x: Tensor,\n attn_mask: Optional[Tensor],\n key_padding_mask: Optional[Tensor],\n cache=None,\n ) -> Tensor:\n x = self.self_attn(\n x,\n x,\n x,\n attn_mask=attn_mask,\n key_padding_mask=key_padding_mask,\n need_weights=False,\n cache=cache,\n )\n return self.dropout1(x)\n\n def _ff_block(self, x: Tensor) -> Tensor:\n x = self.linear2(self.dropout(self.activation(self.linear1(x))))\n return self.dropout2(x)\n\n\nclass AdaptiveLayerNorm(nn.Module):\n r\"\"\"Adaptive Layer Normalization\"\"\"\n\n def __init__(self, d_model, norm) -> None:\n super(AdaptiveLayerNorm, self).__init__()\n self.project_layer = nn.Linear(d_model, 2 * d_model)\n self.norm = norm\n self.d_model = d_model\n self.eps = self.norm.eps\n\n def forward(self, input: Tensor, embedding: Tensor = None) -> Tensor:\n if isinstance(input, tuple):\n input, embedding = input\n weight, bias = torch.split(\n self.project_layer(embedding),\n split_size_or_sections=self.d_model,\n dim=-1,\n )\n return (weight * self.norm(input) + bias, embedding)\n\n weight, bias = torch.split(\n self.project_layer(embedding),\n split_size_or_sections=self.d_model,\n dim=-1,\n )\n return weight * self.norm(input) + bias\n\n\ndef _get_clones(module, N):\n return nn.ModuleList([copy.deepcopy(module) for i in range(N)])\n"
},
{
"path": "GPT_SoVITS/AR/text_processing/__init__.py",
"content": ""
},
{
"path": "GPT_SoVITS/AR/text_processing/phonemizer.py",
"content": "# modified from https://github.com/yangdongchao/SoundStorm/blob/master/soundstorm/s1/AR/text_processing/phonemizer.py\n# reference: https://github.com/lifeiteng/vall-e\nimport itertools\nimport re\nfrom typing import Dict\nfrom typing import List\n\nimport regex\nfrom gruut import sentences\nfrom gruut.const import Sentence\nfrom gruut.const import Word\nfrom AR.text_processing.symbols import SYMBOL_TO_ID\n\n\nclass GruutPhonemizer:\n def __init__(self, language: str):\n self._phonemizer = sentences\n self.lang = language\n self.symbol_to_id = SYMBOL_TO_ID\n self._special_cases_dict: Dict[str] = {\n r\"\\.\\.\\.\": \"... \",\n \";\": \"; \",\n \":\": \": \",\n \",\": \", \",\n r\"\\.\": \". \",\n \"!\": \"! \",\n r\"\\?\": \"? \",\n \"—\": \"—\",\n \"…\": \"… \",\n \"«\": \"«\",\n \"»\": \"»\",\n }\n self._punctuation_regexp: str = rf\"([{''.join(self._special_cases_dict.keys())}])\"\n\n def _normalize_punctuation(self, text: str) -> str:\n text = regex.sub(rf\"\\pZ+{self._punctuation_regexp}\", r\"\\1\", text)\n text = regex.sub(rf\"{self._punctuation_regexp}(\\pL)\", r\"\\1 \\2\", text)\n text = regex.sub(r\"\\pZ+\", r\" \", text)\n return text.strip()\n\n def _convert_punctuation(self, word: Word) -> str:\n if not word.phonemes:\n return \"\"\n if word.phonemes[0] in [\"‖\", \"|\"]:\n return word.text.strip()\n\n phonemes = \"\".join(word.phonemes)\n # remove modifier characters ˈˌː with regex\n phonemes = re.sub(r\"[ˈˌː͡]\", \"\", phonemes)\n return phonemes.strip()\n\n def phonemize(self, text: str, espeak: bool = False) -> str:\n text_to_phonemize: str = self._normalize_punctuation(text)\n sents: List[Sentence] = [sent for sent in self._phonemizer(text_to_phonemize, lang=\"en-us\", espeak=espeak)]\n words: List[str] = [self._convert_punctuation(word) for word in itertools.chain(*sents)]\n return \" \".join(words)\n\n def transform(self, phonemes):\n # convert phonemes to ids\n # dictionary is in symbols.py\n return [self.symbol_to_id[p] for p in phonemes if p in self.symbol_to_id.keys()]\n\n\nif __name__ == \"__main__\":\n phonemizer = GruutPhonemizer(\"en-us\")\n # text -> IPA\n phonemes = phonemizer.phonemize(\"Hello, wor-ld ?\")\n print(\"phonemes:\", phonemes)\n print(\"len(phonemes):\", len(phonemes))\n phoneme_ids = phonemizer.transform(phonemes)\n print(\"phoneme_ids:\", phoneme_ids)\n print(\"len(phoneme_ids):\", len(phoneme_ids))\n"
},
{
"path": "GPT_SoVITS/AR/text_processing/symbols.py",
"content": "# modified from https://github.com/yangdongchao/SoundStorm/blob/master/soundstorm/s1/AR/text_processing/symbols.py\n# reference: https://github.com/lifeiteng/vall-e\nPAD = \"_\"\nPUNCTUATION = ';:,.!?¡¿—…\"«»“” '\nLETTERS = \"ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz\"\nIPA_LETTERS = (\n \"ɑɐɒæɓʙβɔɕçɗɖðʤəɘɚɛɜɝɞɟʄɡɠɢʛɦɧħɥʜɨɪʝɭɬɫɮʟɱɯɰŋɳɲɴøɵɸθœɶʘɹɺɾɻʀʁɽʂʃʈʧʉʊʋⱱʌɣɤʍχʎʏʑʐʒʔʡʕʢǀǁǂǃˈˌːˑʼʴʰʱʲʷˠˤ˞↓↑→↗↘'̩'ᵻ\"\n)\nSYMBOLS = [PAD] + list(PUNCTUATION) + list(LETTERS) + list(IPA_LETTERS)\nSPACE_ID = SYMBOLS.index(\" \")\nSYMBOL_TO_ID = {s: i for i, s in enumerate(SYMBOLS)}\nID_TO_SYMBOL = {i: s for i, s in enumerate(SYMBOLS)}\n"
},
{
"path": "GPT_SoVITS/AR/utils/__init__.py",
"content": "import re\n\n\ndef str2bool(str):\n return True if str.lower() == \"true\" else False\n\n\ndef get_newest_ckpt(string_list):\n # 定义一个正则表达式模式，用于匹配字符串中的数字\n pattern = r\"epoch=(\\d+)-step=(\\d+)\\.ckpt\"\n\n # 使用正则表达式提取每个字符串中的数字信息，并创建一个包含元组的列表\n extracted_info = []\n for string in string_list:\n match = re.match(pattern, string)\n if match:\n epoch = int(match.group(1))\n step = int(match.group(2))\n extracted_info.append((epoch, step, string))\n # 按照 epoch 后面的数字和 step 后面的数字进行排序\n sorted_info = sorted(extracted_info, key=lambda x: (x[0], x[1]), reverse=True)\n # 获取最新的 ckpt 文件名\n newest_ckpt = sorted_info[0][2]\n return newest_ckpt\n\n\n# 文本存在且不为空时 return True\ndef check_txt_file(file_path):\n try:\n with open(file_path, \"r\") as file:\n text = file.readline().strip()\n assert text.strip() != \"\"\n return text\n except Exception:\n return False\n return False\n"
},
{
"path": "GPT_SoVITS/AR/utils/initialize.py",
"content": "#!/usr/bin/env python3\n\"\"\"Initialize modules for espnet2 neural networks.\"\"\"\n\nimport torch\nfrom typeguard import check_argument_types\n\n\ndef initialize(model: torch.nn.Module, init: str):\n \"\"\"Initialize weights of a neural network module.\n\n Parameters are initialized using the given method or distribution.\n\n Custom initialization routines can be implemented into submodules\n as function `espnet_initialization_fn` within the custom module.\n\n Args:\n model: Target.\n init: Method of initialization.\n \"\"\"\n assert check_argument_types()\n print(\"init with\", init)\n\n # weight init\n for p in model.parameters():\n if p.dim() > 1:\n if init == \"xavier_uniform\":\n torch.nn.init.xavier_uniform_(p.data)\n elif init == \"xavier_normal\":\n torch.nn.init.xavier_normal_(p.data)\n elif init == \"kaiming_uniform\":\n torch.nn.init.kaiming_uniform_(p.data, nonlinearity=\"relu\")\n elif init == \"kaiming_normal\":\n torch.nn.init.kaiming_normal_(p.data, nonlinearity=\"relu\")\n else:\n raise ValueError(\"Unknown initialization: \" + init)\n # bias init\n for name, p in model.named_parameters():\n if \".bias\" in name and p.dim() == 1:\n p.data.zero_()\n"
},
{
"path": "GPT_SoVITS/AR/utils/io.py",
"content": "import sys\n\nimport torch\nimport yaml\n\n\ndef load_yaml_config(path):\n with open(path) as f:\n config = yaml.full_load(f)\n return config\n\n\ndef save_config_to_yaml(config, path):\n assert path.endswith(\".yaml\")\n with open(path, \"w\") as f:\n f.write(yaml.dump(config))\n f.close()\n\n\ndef write_args(args, path):\n args_dict = dict((name, getattr(args, name)) for name in dir(args) if not name.startswith(\"_\"))\n with open(path, \"a\") as args_file:\n args_file.write(\"==> torch version: {}\\n\".format(torch.__version__))\n args_file.write(\"==> cudnn version: {}\\n\".format(torch.backends.cudnn.version()))\n args_file.write(\"==> Cmd:\\n\")\n args_file.write(str(sys.argv))\n args_file.write(\"\\n==> args:\\n\")\n for k, v in sorted(args_dict.items()):\n args_file.write(\" %s: %s\\n\" % (str(k), str(v)))\n args_file.close()\n"
},
{
"path": "GPT_SoVITS/BigVGAN/LICENSE",
"content": "MIT License\n\nCopyright (c) 2024 NVIDIA CORPORATION.\n\nPermission is hereby granted, free of charge, to any person obtaining a copy\nof this software and associated documentation files (the \"Software\"), to deal\nin the Software without restriction, including without limitation the rights\nto use, copy, modify, merge, publish, distribute, sublicense, and/or sell\ncopies of the Software, and to permit persons to whom the Software is\nfurnished to do so, subject to the following conditions:\n\nThe above copyright notice and this permission notice shall be included in all\ncopies or substantial portions of the Software.\n\nTHE SOFTWARE IS PROVIDED \"AS IS\", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR\nIMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,\nFITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE\nAUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER\nLIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,\nOUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE\nSOFTWARE.\n"
},
{
"path": "GPT_SoVITS/BigVGAN/README.md",
"content": "## BigVGAN: A Universal Neural Vocoder with Large-Scale Training\n\n#### Sang-gil Lee, Wei Ping, Boris Ginsburg, Bryan Catanzaro, Sungroh Yoon\n\n[[Paper]](https://arxiv.org/abs/2206.04658) - [[Code]](https://github.com/NVIDIA/BigVGAN) - [[Showcase]](https://bigvgan-demo.github.io/) - [[Project Page]](https://research.nvidia.com/labs/adlr/projects/bigvgan/) - [[Weights]](https://huggingface.co/collections/nvidia/bigvgan-66959df3d97fd7d98d97dc9a) - [[Demo]](https://huggingface.co/spaces/nvidia/BigVGAN)\n\n[](https://paperswithcode.com/sota/speech-synthesis-on-libritts?p=bigvgan-a-universal-neural-vocoder-with-large)\n\n![](\"https://user-images.githubusercontent.com/15963413/218609148-881e39df-33af-4af9-ab95-1427c4ebf062.png\")
\n\nYou can run a local gradio demo using below command:\n\n```python\npip install -r demo/requirements.txt\npython demo/app.py\n```\n\n## Training\n\nCreate symbolic link to the root of the dataset. The codebase uses filelist with the relative path from the dataset. Below are the example commands for LibriTTS dataset:\n\n```shell\ncd filelists/LibriTTS && \\\nln -s /path/to/your/LibriTTS/train-clean-100 train-clean-100 && \\\nln -s /path/to/your/LibriTTS/train-clean-360 train-clean-360 && \\\nln -s /path/to/your/LibriTTS/train-other-500 train-other-500 && \\\nln -s /path/to/your/LibriTTS/dev-clean dev-clean && \\\nln -s /path/to/your/LibriTTS/dev-other dev-other && \\\nln -s /path/to/your/LibriTTS/test-clean test-clean && \\\nln -s /path/to/your/LibriTTS/test-other test-other && \\\ncd ../..\n```\n\nTrain BigVGAN model. Below is an example command for training BigVGAN-v2 using LibriTTS dataset at 24kHz with a full 100-band mel spectrogram as input:\n\n```shell\npython train.py \\\n--config configs/bigvgan_v2_24khz_100band_256x.json \\\n--input_wavs_dir filelists/LibriTTS \\\n--input_training_file filelists/LibriTTS/train-full.txt \\\n--input_validation_file filelists/LibriTTS/val-full.txt \\\n--list_input_unseen_wavs_dir filelists/LibriTTS filelists/LibriTTS \\\n--list_input_unseen_validation_file filelists/LibriTTS/dev-clean.txt filelists/LibriTTS/dev-other.txt \\\n--checkpoint_path exp/bigvgan_v2_24khz_100band_256x\n```\n\n## Synthesis\n\nSynthesize from BigVGAN model. Below is an example command for generating audio from the model.\nIt computes mel spectrograms using wav files from `--input_wavs_dir` and saves the generated audio to `--output_dir`.\n\n```shell\npython inference.py \\\n--checkpoint_file /path/to/your/bigvgan_v2_24khz_100band_256x/bigvgan_generator.pt \\\n--input_wavs_dir /path/to/your/input_wav \\\n--output_dir /path/to/your/output_wav\n```\n\n`inference_e2e.py` supports synthesis directly from the mel spectrogram saved in `.npy` format, with shapes `[1, channel, frame]` or `[channel, frame]`.\nIt loads mel spectrograms from `--input_mels_dir` and saves the generated audio to `--output_dir`.\n\nMake sure that the STFT hyperparameters for mel spectrogram are the same as the model, which are defined in `config.json` of the corresponding model.\n\n```shell\npython inference_e2e.py \\\n--checkpoint_file /path/to/your/bigvgan_v2_24khz_100band_256x/bigvgan_generator.pt \\\n--input_mels_dir /path/to/your/input_mel \\\n--output_dir /path/to/your/output_wav\n```\n\n## Using Custom CUDA Kernel for Synthesis\n\nYou can apply the fast CUDA inference kernel by using a parameter `use_cuda_kernel` when instantiating BigVGAN:\n\n```python\ngenerator = BigVGAN(h, use_cuda_kernel=True)\n```\n\nYou can also pass `--use_cuda_kernel` to `inference.py` and `inference_e2e.py` to enable this feature.\n\nWhen applied for the first time, it builds the kernel using `nvcc` and `ninja`. If the build succeeds, the kernel is saved to `alias_free_activation/cuda/build` and the model automatically loads the kernel. The codebase has been tested using CUDA `12.1`.\n\nPlease make sure that both are installed in your system and `nvcc` installed in your system matches the version your PyTorch build is using.\n\nWe recommend running `test_cuda_vs_torch_model.py` first to build and check the correctness of the CUDA kernel. See below example command and its output, where it returns `[Success] test CUDA fused vs. plain torch BigVGAN inference`:\n\n```python\npython tests/test_cuda_vs_torch_model.py \\\n--checkpoint_file /path/to/your/bigvgan_generator.pt\n```\n\n```shell\nloading plain Pytorch BigVGAN\n...\nloading CUDA kernel BigVGAN with auto-build\nDetected CUDA files, patching ldflags\nEmitting ninja build file /path/to/your/BigVGAN/alias_free_activation/cuda/build/build.ninja..\nBuilding extension module anti_alias_activation_cuda...\n...\nLoading extension module anti_alias_activation_cuda...\n...\nLoading '/path/to/your/bigvgan_generator.pt'\n...\n[Success] test CUDA fused vs. plain torch BigVGAN inference\n > mean_difference=0.0007238413265440613\n...\n```\n\nIf you see `[Fail] test CUDA fused vs. plain torch BigVGAN inference`, it means that the CUDA kernel inference is incorrect. Please check if `nvcc` installed in your system is compatible with your PyTorch version.\n\n## Pretrained Models\n\nWe provide the [pretrained models on Hugging Face Collections](https://huggingface.co/collections/nvidia/bigvgan-66959df3d97fd7d98d97dc9a).\nOne can download the checkpoints of the generator weight (named `bigvgan_generator.pt`) and its discriminator/optimizer states (named `bigvgan_discriminator_optimizer.pt`) within the listed model repositories.\n\n| Model Name | Sampling Rate | Mel band | fmax | Upsampling Ratio | Params | Dataset | Steps | Fine-Tuned |\n|:--------------------------------------------------------------------------------------------------------:|:-------------:|:--------:|:-----:|:----------------:|:------:|:--------------------------:|:-----:|:----------:|\n| [bigvgan_v2_44khz_128band_512x](https://huggingface.co/nvidia/bigvgan_v2_44khz_128band_512x) | 44 kHz | 128 | 22050 | 512 | 122M | Large-scale Compilation | 5M | No |\n| [bigvgan_v2_44khz_128band_256x](https://huggingface.co/nvidia/bigvgan_v2_44khz_128band_256x) | 44 kHz | 128 | 22050 | 256 | 112M | Large-scale Compilation | 5M | No |\n| [bigvgan_v2_24khz_100band_256x](https://huggingface.co/nvidia/bigvgan_v2_24khz_100band_256x) | 24 kHz | 100 | 12000 | 256 | 112M | Large-scale Compilation | 5M | No |\n| [bigvgan_v2_22khz_80band_256x](https://huggingface.co/nvidia/bigvgan_v2_22khz_80band_256x) | 22 kHz | 80 | 11025 | 256 | 112M | Large-scale Compilation | 5M | No |\n| [bigvgan_v2_22khz_80band_fmax8k_256x](https://huggingface.co/nvidia/bigvgan_v2_22khz_80band_fmax8k_256x) | 22 kHz | 80 | 8000 | 256 | 112M | Large-scale Compilation | 5M | No |\n| [bigvgan_24khz_100band](https://huggingface.co/nvidia/bigvgan_24khz_100band) | 24 kHz | 100 | 12000 | 256 | 112M | LibriTTS | 5M | No |\n| [bigvgan_base_24khz_100band](https://huggingface.co/nvidia/bigvgan_base_24khz_100band) | 24 kHz | 100 | 12000 | 256 | 14M | LibriTTS | 5M | No |\n| [bigvgan_22khz_80band](https://huggingface.co/nvidia/bigvgan_22khz_80band) | 22 kHz | 80 | 8000 | 256 | 112M | LibriTTS + VCTK + LJSpeech | 5M | No |\n| [bigvgan_base_22khz_80band](https://huggingface.co/nvidia/bigvgan_base_22khz_80band) | 22 kHz | 80 | 8000 | 256 | 14M | LibriTTS + VCTK + LJSpeech | 5M | No |\n\nThe paper results are based on the original 24kHz BigVGAN models (`bigvgan_24khz_100band` and `bigvgan_base_24khz_100band`) trained on LibriTTS dataset.\nWe also provide 22kHz BigVGAN models with band-limited setup (i.e., fmax=8000) for TTS applications.\nNote that the checkpoints use `snakebeta` activation with log scale parameterization, which have the best overall quality.\n\nYou can fine-tune the models by:\n\n1. downloading the checkpoints (both the generator weight and its discriminator/optimizer states)\n2. resuming training using your audio dataset by specifying `--checkpoint_path` that includes the checkpoints when launching `train.py`\n\n## Training Details of BigVGAN-v2\n\nComapred to the original BigVGAN, the pretrained checkpoints of BigVGAN-v2 used `batch_size=32` with a longer `segment_size=65536` and are trained using 8 A100 GPUs.\n\nNote that the BigVGAN-v2 `json` config files in `./configs` use `batch_size=4` as default to fit in a single A100 GPU for training. You can fine-tune the models adjusting `batch_size` depending on your GPUs.\n\nWhen training BigVGAN-v2 from scratch with small batch size, it can potentially encounter the early divergence problem mentioned in the paper. In such case, we recommend lowering the `clip_grad_norm` value (e.g. `100`) for the early training iterations (e.g. 20000 steps) and increase the value to the default `500`.\n\n## Evaluation Results of BigVGAN-v2\n\nBelow are the objective results of the 24kHz model (`bigvgan_v2_24khz_100band_256x`) obtained from the LibriTTS `dev` sets. BigVGAN-v2 shows noticeable improvements of the metrics. The model also exhibits reduced perceptual artifacts, especially for non-speech audio.\n\n| Model | Dataset | Steps | PESQ(↑) | M-STFT(↓) | MCD(↓) | Periodicity(↓) | V/UV F1(↑) |\n|:----------:|:-----------------------:|:-----:|:---------:|:----------:|:----------:|:--------------:|:----------:|\n| BigVGAN | LibriTTS | 1M | 4.027 | 0.7997 | 0.3745 | 0.1018 | 0.9598 |\n| BigVGAN | LibriTTS | 5M | 4.256 | 0.7409 | 0.2988 | 0.0809 | 0.9698 |\n| BigVGAN-v2 | Large-scale Compilation | 3M | 4.359 | 0.7134 | 0.3060 | 0.0621 | 0.9777 |\n| BigVGAN-v2 | Large-scale Compilation | 5M | **4.362** | **0.7026** | **0.2903** | **0.0593** | **0.9793** |\n\n## Speed Benchmark\n\nBelow are the speed and VRAM usage benchmark results of BigVGAN from `tests/test_cuda_vs_torch_model.py`, using `bigvgan_v2_24khz_100band_256x` as a reference model.\n\n| GPU | num_mel_frame | use_cuda_kernel | Speed (kHz) | Real-time Factor | VRAM (GB) |\n|:--------------------------:|:-------------:|:---------------:|:-----------:|:----------------:|:---------:|\n| NVIDIA A100 | 256 | False | 1672.1 | 69.7x | 1.3 |\n| | | True | 3916.5 | 163.2x | 1.3 |\n| | 2048 | False | 1899.6 | 79.2x | 1.7 |\n| | | True | 5330.1 | 222.1x | 1.7 |\n| | 16384 | False | 1973.8 | 82.2x | 5.0 |\n| | | True | 5761.7 | 240.1x | 4.4 |\n| NVIDIA GeForce RTX 3080 | 256 | False | 841.1 | 35.0x | 1.3 |\n| | | True | 1598.1 | 66.6x | 1.3 |\n| | 2048 | False | 929.9 | 38.7x | 1.7 |\n| | | True | 1971.3 | 82.1x | 1.6 |\n| | 16384 | False | 943.4 | 39.3x | 5.0 |\n| | | True | 2026.5 | 84.4x | 3.9 |\n| NVIDIA GeForce RTX 2080 Ti | 256 | False | 515.6 | 21.5x | 1.3 |\n| | | True | 811.3 | 33.8x | 1.3 |\n| | 2048 | False | 576.5 | 24.0x | 1.7 |\n| | | True | 1023.0 | 42.6x | 1.5 |\n| | 16384 | False | 589.4 | 24.6x | 5.0 |\n| | | True | 1068.1 | 44.5x | 3.2 |\n\n## Acknowledgements\n\nWe thank Vijay Anand Korthikanti and Kevin J. Shih for their generous support in implementing the CUDA kernel for inference.\n\n## References\n\n- [HiFi-GAN](https://github.com/jik876/hifi-gan) (for generator and multi-period discriminator)\n- [Snake](https://github.com/EdwardDixon/snake) (for periodic activation)\n- [Alias-free-torch](https://github.com/junjun3518/alias-free-torch) (for anti-aliasing)\n- [Julius](https://github.com/adefossez/julius) (for low-pass filter)\n- [UnivNet](https://github.com/mindslab-ai/univnet) (for multi-resolution discriminator)\n- [descript-audio-codec](https://github.com/descriptinc/descript-audio-codec) and [vocos](https://github.com/gemelo-ai/vocos) (for multi-band multi-scale STFT discriminator and multi-scale mel spectrogram loss)\n- [Amphion](https://github.com/open-mmlab/Amphion) (for multi-scale sub-band CQT discriminator)\n"
},
{
"path": "GPT_SoVITS/BigVGAN/activations.py",
"content": "# Implementation adapted from https://github.com/EdwardDixon/snake under the MIT license.\n# LICENSE is in incl_licenses directory.\n\nimport torch\nfrom torch import nn, sin, pow\nfrom torch.nn import Parameter\n\n\nclass Snake(nn.Module):\n \"\"\"\n Implementation of a sine-based periodic activation function\n Shape:\n - Input: (B, C, T)\n - Output: (B, C, T), same shape as the input\n Parameters:\n - alpha - trainable parameter\n References:\n - This activation function is from this paper by Liu Ziyin, Tilman Hartwig, Masahito Ueda:\n https://arxiv.org/abs/2006.08195\n Examples:\n >>> a1 = snake(256)\n >>> x = torch.randn(256)\n >>> x = a1(x)\n \"\"\"\n\n def __init__(self, in_features, alpha=1.0, alpha_trainable=True, alpha_logscale=False):\n \"\"\"\n Initialization.\n INPUT:\n - in_features: shape of the input\n - alpha: trainable parameter\n alpha is initialized to 1 by default, higher values = higher-frequency.\n alpha will be trained along with the rest of your model.\n \"\"\"\n super(Snake, self).__init__()\n self.in_features = in_features\n\n # Initialize alpha\n self.alpha_logscale = alpha_logscale\n if self.alpha_logscale: # Log scale alphas initialized to zeros\n self.alpha = Parameter(torch.zeros(in_features) * alpha)\n else: # Linear scale alphas initialized to ones\n self.alpha = Parameter(torch.ones(in_features) * alpha)\n\n self.alpha.requires_grad = alpha_trainable\n\n self.no_div_by_zero = 0.000000001\n\n def forward(self, x):\n \"\"\"\n Forward pass of the function.\n Applies the function to the input elementwise.\n Snake ∶= x + 1/a * sin^2 (xa)\n \"\"\"\n alpha = self.alpha.unsqueeze(0).unsqueeze(-1) # Line up with x to [B, C, T]\n if self.alpha_logscale:\n alpha = torch.exp(alpha)\n x = x + (1.0 / (alpha + self.no_div_by_zero)) * pow(sin(x * alpha), 2)\n\n return x\n\n\nclass SnakeBeta(nn.Module):\n \"\"\"\n A modified Snake function which uses separate parameters for the magnitude of the periodic components\n Shape:\n - Input: (B, C, T)\n - Output: (B, C, T), same shape as the input\n Parameters:\n - alpha - trainable parameter that controls frequency\n - beta - trainable parameter that controls magnitude\n References:\n - This activation function is a modified version based on this paper by Liu Ziyin, Tilman Hartwig, Masahito Ueda:\n https://arxiv.org/abs/2006.08195\n Examples:\n >>> a1 = snakebeta(256)\n >>> x = torch.randn(256)\n >>> x = a1(x)\n \"\"\"\n\n def __init__(self, in_features, alpha=1.0, alpha_trainable=True, alpha_logscale=False):\n \"\"\"\n Initialization.\n INPUT:\n - in_features: shape of the input\n - alpha - trainable parameter that controls frequency\n - beta - trainable parameter that controls magnitude\n alpha is initialized to 1 by default, higher values = higher-frequency.\n beta is initialized to 1 by default, higher values = higher-magnitude.\n alpha will be trained along with the rest of your model.\n \"\"\"\n super(SnakeBeta, self).__init__()\n self.in_features = in_features\n\n # Initialize alpha\n self.alpha_logscale = alpha_logscale\n if self.alpha_logscale: # Log scale alphas initialized to zeros\n self.alpha = Parameter(torch.zeros(in_features) * alpha)\n self.beta = Parameter(torch.zeros(in_features) * alpha)\n else: # Linear scale alphas initialized to ones\n self.alpha = Parameter(torch.ones(in_features) * alpha)\n self.beta = Parameter(torch.ones(in_features) * alpha)\n\n self.alpha.requires_grad = alpha_trainable\n self.beta.requires_grad = alpha_trainable\n\n self.no_div_by_zero = 0.000000001\n\n def forward(self, x):\n \"\"\"\n Forward pass of the function.\n Applies the function to the input elementwise.\n SnakeBeta ∶= x + 1/b * sin^2 (xa)\n \"\"\"\n alpha = self.alpha.unsqueeze(0).unsqueeze(-1) # Line up with x to [B, C, T]\n beta = self.beta.unsqueeze(0).unsqueeze(-1)\n if self.alpha_logscale:\n alpha = torch.exp(alpha)\n beta = torch.exp(beta)\n x = x + (1.0 / (beta + self.no_div_by_zero)) * pow(sin(x * alpha), 2)\n\n return x\n"
},
{
"path": "GPT_SoVITS/BigVGAN/alias_free_activation/cuda/__init__.py",
"content": ""
},
{
"path": "GPT_SoVITS/BigVGAN/alias_free_activation/cuda/activation1d.py",
"content": "# Copyright (c) 2024 NVIDIA CORPORATION.\n# Licensed under the MIT license.\n\nimport torch\nimport torch.nn as nn\nfrom alias_free_activation.torch.resample import UpSample1d, DownSample1d\n\n# load fused CUDA kernel: this enables importing anti_alias_activation_cuda\nfrom alias_free_activation.cuda import load\n\nanti_alias_activation_cuda = load.load()\n\n\nclass FusedAntiAliasActivation(torch.autograd.Function):\n \"\"\"\n Assumes filter size 12, replication padding on upsampling/downsampling, and logscale alpha/beta parameters as inputs.\n The hyperparameters are hard-coded in the kernel to maximize speed.\n NOTE: The fused kenrel is incorrect for Activation1d with different hyperparameters.\n \"\"\"\n\n @staticmethod\n def forward(ctx, inputs, up_ftr, down_ftr, alpha, beta):\n activation_results = anti_alias_activation_cuda.forward(inputs, up_ftr, down_ftr, alpha, beta)\n\n return activation_results\n\n @staticmethod\n def backward(ctx, output_grads):\n raise NotImplementedError\n return output_grads, None, None\n\n\nclass Activation1d(nn.Module):\n def __init__(\n self,\n activation,\n up_ratio: int = 2,\n down_ratio: int = 2,\n up_kernel_size: int = 12,\n down_kernel_size: int = 12,\n fused: bool = True,\n ):\n super().__init__()\n self.up_ratio = up_ratio\n self.down_ratio = down_ratio\n self.act = activation\n self.upsample = UpSample1d(up_ratio, up_kernel_size)\n self.downsample = DownSample1d(down_ratio, down_kernel_size)\n\n self.fused = fused # Whether to use fused CUDA kernel or not\n\n def forward(self, x):\n if not self.fused:\n x = self.upsample(x)\n x = self.act(x)\n x = self.downsample(x)\n return x\n else:\n if self.act.__class__.__name__ == \"Snake\":\n beta = self.act.alpha.data # Snake uses same params for alpha and beta\n else:\n beta = self.act.beta.data # Snakebeta uses different params for alpha and beta\n alpha = self.act.alpha.data\n if not self.act.alpha_logscale: # Exp baked into cuda kernel, cancel it out with a log\n alpha = torch.log(alpha)\n beta = torch.log(beta)\n\n x = FusedAntiAliasActivation.apply(x, self.upsample.filter, self.downsample.lowpass.filter, alpha, beta)\n return x\n"
},
{
"path": "GPT_SoVITS/BigVGAN/alias_free_activation/cuda/anti_alias_activation.cpp",
"content": "/* coding=utf-8\n * Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved.\n *\n * Licensed under the Apache License, Version 2.0 (the \"License\");\n * you may not use this file except in compliance with the License.\n * You may obtain a copy of the License at\n *\n * http://www.apache.org/licenses/LICENSE-2.0\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */\n\n #include