mirror of
https://github.com/RVC-Boss/GPT-SoVITS.git
synced 2026-08-13 01:53:43 +08:00
Compare commits
79
Commits
20250606v2pro
...
main
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
d523079fc0 | ||
|
|
be6a4f1e9d | ||
|
|
551918539c | ||
|
|
bf81cdb14a | ||
|
|
b2cff0cd0a | ||
|
|
08d627c333 | ||
|
|
6d95b559e8 | ||
|
|
ea2d2a8166 | ||
|
|
d9f03dad3e | ||
|
|
647935357a | ||
|
|
02425ea256 | ||
|
|
938f05fce8 | ||
|
|
445d18ccce | ||
|
|
00ce973412 | ||
|
|
14191901cd | ||
|
|
780383d5bd | ||
|
|
ba8de9b760 | ||
|
|
2d9193b0d3 | ||
|
|
9986880b3f | ||
|
|
c767f0b83b | ||
|
|
9080a967d5 | ||
|
|
51df9f7384 | ||
|
|
bfca0f6b2d | ||
|
|
abe984395c | ||
|
|
cc89c3660e | ||
|
|
36b3231c6f | ||
|
|
9ec3a60f30 | ||
|
|
fc533b6fb7 | ||
|
|
857799276c | ||
|
|
92d2d337fd | ||
|
|
6fb441f65e | ||
|
|
c85c54eca9 | ||
|
|
cb00840c4e | ||
|
|
60a4a214af | ||
|
|
6375bbe316 | ||
|
|
e00ca92140 | ||
|
|
92ab59c553 | ||
|
|
11aa78bd9b | ||
|
|
fdf794e31d | ||
|
|
0be59c8043 | ||
|
|
b5a67e6247 | ||
|
|
b9211657d8 | ||
|
|
cefafee32c | ||
|
|
2d09bbe63a | ||
|
|
4d8ebf8523 | ||
|
|
e476b01f30 | ||
|
|
42586e20f7 | ||
|
|
85035f7ac0 | ||
|
|
706bec74f8 | ||
|
|
ec1218893e | ||
|
|
fec515dcce | ||
|
|
426e1a2bb4 | ||
|
|
4e3c69043c | ||
|
|
e63e0901fd | ||
|
|
97e37c74d8 | ||
|
|
3a75f5023f | ||
|
|
0899b7e432 | ||
|
|
8c579d46dd | ||
|
|
6df61f58e4 | ||
|
|
90ebefa78f | ||
|
|
4839e82148 | ||
|
|
37f5abfcb4 | ||
|
|
4987df5a71 | ||
|
|
d46c069e52 | ||
|
|
6fdc67ca83 | ||
|
|
7dec5f5bb0 | ||
|
|
1a9b8854ee | ||
|
|
5c91e66d2e | ||
|
|
ed89a02337 | ||
|
|
cd6de7398e | ||
|
|
dd2b9253aa | ||
|
|
29165eb02e | ||
|
|
746cb536c6 | ||
|
|
0d2f273402 | ||
|
|
d39836b8fa | ||
|
|
2c0436b9ce | ||
|
|
8056efe4ab | ||
|
|
d6b78c927a | ||
|
|
74e79ae6d6 |
@@ -21,6 +21,7 @@ jobs:
|
||||
build-amd64:
|
||||
needs: generate-meta
|
||||
runs-on: ubuntu-22.04
|
||||
environment: Docker
|
||||
strategy:
|
||||
matrix:
|
||||
include:
|
||||
@@ -86,7 +87,7 @@ jobs:
|
||||
- name: Log in to Docker Hub
|
||||
uses: docker/login-action@v3
|
||||
with:
|
||||
username: ${{ secrets.DOCKER_HUB_USERNAME }}
|
||||
username: ${{ vars.DOCKER_HUB_USERNAME }}
|
||||
password: ${{ secrets.DOCKER_HUB_PASSWORD }}
|
||||
|
||||
- name: Build and Push Docker Image (amd64)
|
||||
@@ -108,6 +109,7 @@ jobs:
|
||||
build-arm64:
|
||||
needs: generate-meta
|
||||
runs-on: ubuntu-22.04-arm
|
||||
environment: Docker
|
||||
strategy:
|
||||
matrix:
|
||||
include:
|
||||
@@ -173,7 +175,7 @@ jobs:
|
||||
- name: Log in to Docker Hub
|
||||
uses: docker/login-action@v3
|
||||
with:
|
||||
username: ${{ secrets.DOCKER_HUB_USERNAME }}
|
||||
username: ${{ vars.DOCKER_HUB_USERNAME }}
|
||||
password: ${{ secrets.DOCKER_HUB_PASSWORD }}
|
||||
|
||||
- name: Build and Push Docker Image (arm64)
|
||||
@@ -206,6 +208,7 @@ jobs:
|
||||
- tag_prefix: cu126
|
||||
- tag_prefix: cu128-lite
|
||||
- tag_prefix: cu128
|
||||
environment: Docker
|
||||
|
||||
steps:
|
||||
- name: Set up Docker Buildx
|
||||
@@ -214,7 +217,7 @@ jobs:
|
||||
- name: Log in to Docker Hub
|
||||
uses: docker/login-action@v3
|
||||
with:
|
||||
username: ${{ secrets.DOCKER_HUB_USERNAME }}
|
||||
username: ${{ vars.DOCKER_HUB_USERNAME }}
|
||||
password: ${{ secrets.DOCKER_HUB_PASSWORD }}
|
||||
|
||||
- name: Merge amd64 and arm64 into multi-arch image
|
||||
@@ -223,17 +226,17 @@ jobs:
|
||||
TAG_PREFIX=${{ matrix.tag_prefix }}
|
||||
|
||||
docker buildx imagetools create \
|
||||
--tag ${{ secrets.DOCKER_HUB_USERNAME }}/gpt-sovits:${TAG_PREFIX}-${DATE_TAG} \
|
||||
${{ secrets.DOCKER_HUB_USERNAME }}/gpt-sovits:${TAG_PREFIX}-${DATE_TAG}-amd64 \
|
||||
${{ secrets.DOCKER_HUB_USERNAME }}/gpt-sovits:${TAG_PREFIX}-${DATE_TAG}-arm64
|
||||
--tag ${{ vars.DOCKER_HUB_USERNAME }}/gpt-sovits:${TAG_PREFIX}-${DATE_TAG} \
|
||||
${{ vars.DOCKER_HUB_USERNAME }}/gpt-sovits:${TAG_PREFIX}-${DATE_TAG}-amd64 \
|
||||
${{ vars.DOCKER_HUB_USERNAME }}/gpt-sovits:${TAG_PREFIX}-${DATE_TAG}-arm64
|
||||
|
||||
docker buildx imagetools create \
|
||||
--tag ${{ secrets.DOCKER_HUB_USERNAME }}/gpt-sovits:latest-${TAG_PREFIX} \
|
||||
${{ secrets.DOCKER_HUB_USERNAME }}/gpt-sovits:latest-${TAG_PREFIX}-amd64 \
|
||||
${{ secrets.DOCKER_HUB_USERNAME }}/gpt-sovits:latest-${TAG_PREFIX}-arm64
|
||||
--tag ${{ vars.DOCKER_HUB_USERNAME }}/gpt-sovits:latest-${TAG_PREFIX} \
|
||||
${{ vars.DOCKER_HUB_USERNAME }}/gpt-sovits:latest-${TAG_PREFIX}-amd64 \
|
||||
${{ vars.DOCKER_HUB_USERNAME }}/gpt-sovits:latest-${TAG_PREFIX}-arm64
|
||||
- name: Delete old platform-specific tags via Docker Hub API
|
||||
env:
|
||||
DOCKER_HUB_USERNAME: ${{ secrets.DOCKER_HUB_USERNAME }}
|
||||
DOCKER_HUB_USERNAME: ${{ vars.DOCKER_HUB_USERNAME }}
|
||||
DOCKER_HUB_TOKEN: ${{ secrets.DOCKER_HUB_PASSWORD }}
|
||||
TAG_PREFIX: ${{ matrix.tag_prefix }}
|
||||
DATE_TAG: ${{ needs.generate-meta.outputs.tag }}
|
||||
@@ -258,6 +261,7 @@ jobs:
|
||||
runs-on: ubuntu-latest
|
||||
needs:
|
||||
- merge-and-clean
|
||||
environment: Docker
|
||||
steps:
|
||||
- name: Set up Docker Buildx
|
||||
uses: docker/setup-buildx-action@v3
|
||||
@@ -265,12 +269,12 @@ jobs:
|
||||
- name: Log in to Docker Hub
|
||||
uses: docker/login-action@v3
|
||||
with:
|
||||
username: ${{ secrets.DOCKER_HUB_USERNAME }}
|
||||
username: ${{ vars.DOCKER_HUB_USERNAME }}
|
||||
password: ${{ secrets.DOCKER_HUB_PASSWORD }}
|
||||
|
||||
- name: Create Default Tag
|
||||
run: |
|
||||
docker buildx imagetools create \
|
||||
--tag ${{ secrets.DOCKER_HUB_USERNAME }}/gpt-sovits:latest \
|
||||
${{ secrets.DOCKER_HUB_USERNAME }}/gpt-sovits:latest-cu126-lite
|
||||
--tag ${{ vars.DOCKER_HUB_USERNAME }}/gpt-sovits:latest \
|
||||
${{ vars.DOCKER_HUB_USERNAME }}/gpt-sovits:latest-cu126-lite
|
||||
|
||||
@@ -12,4 +12,4 @@ repos:
|
||||
# Run the formatter.
|
||||
- id: ruff-format
|
||||
types_or: [ python, pyi ]
|
||||
args: [ --line-length, "120", --target-version, "py310" ]
|
||||
args: [ --line-length, "120", --target-version, "py311" ]
|
||||
|
||||
@@ -8,7 +8,7 @@ cd .. || exit 1
|
||||
|
||||
set -e
|
||||
|
||||
source "$HOME/miniconda3/etc/profile.d/conda.sh"
|
||||
source "$HOME/conda/etc/profile.d/conda.sh"
|
||||
|
||||
mkdir -p GPT_SoVITS
|
||||
|
||||
@@ -18,7 +18,7 @@ ln -s /workspace/models/pretrained_models /workspace/GPT-SoVITS/GPT_SoVITS/pretr
|
||||
|
||||
ln -s /workspace/models/G2PWModel /workspace/GPT-SoVITS/GPT_SoVITS/text/G2PWModel
|
||||
|
||||
bash install.sh --device "CU${CUDA_VERSION//./}" --source HF
|
||||
TERM=dumb bash install.sh --device "CU${CUDA_VERSION//./}" --source HF
|
||||
|
||||
pip cache purge
|
||||
|
||||
@@ -26,8 +26,8 @@ pip show torch
|
||||
|
||||
rm -rf /tmp/* /var/tmp/*
|
||||
|
||||
rm -rf "$HOME/miniconda3/pkgs"
|
||||
rm -rf "$HOME/conda/pkgs"
|
||||
|
||||
mkdir -p "$HOME/miniconda3/pkgs"
|
||||
mkdir -p "$HOME/conda/pkgs"
|
||||
|
||||
rm -rf /root/.conda /root/.cache
|
||||
|
||||
@@ -1,70 +0,0 @@
|
||||
#!/bin/bash
|
||||
|
||||
set -e
|
||||
|
||||
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" &>/dev/null && pwd)"
|
||||
|
||||
cd "$SCRIPT_DIR" || exit 1
|
||||
|
||||
cd .. || exit 1
|
||||
|
||||
if [ -d "$HOME/miniconda3" ]; then
|
||||
exit 0
|
||||
fi
|
||||
|
||||
WORKFLOW=${WORKFLOW:-"false"}
|
||||
TARGETPLATFORM=${TARGETPLATFORM:-"linux/amd64"}
|
||||
|
||||
if [ "$WORKFLOW" = "true" ]; then
|
||||
WGET_CMD=(wget -nv --tries=25 --wait=5 --read-timeout=40 --retry-on-http-error=404)
|
||||
else
|
||||
WGET_CMD=(wget --tries=25 --wait=5 --read-timeout=40 --retry-on-http-error=404)
|
||||
fi
|
||||
|
||||
if [ "$TARGETPLATFORM" = "linux/amd64" ]; then
|
||||
"${WGET_CMD[@]}" -O miniconda.sh https://repo.anaconda.com/miniconda/Miniconda3-py311_25.3.1-1-Linux-x86_64.sh
|
||||
elif [ "$TARGETPLATFORM" = "linux/arm64" ]; then
|
||||
"${WGET_CMD[@]}" -O miniconda.sh https://repo.anaconda.com/miniconda/Miniconda3-py311_25.3.1-1-Linux-aarch64.sh
|
||||
else
|
||||
exit 1
|
||||
fi
|
||||
|
||||
LOG_PATH="/tmp/miniconda-install.log"
|
||||
|
||||
bash miniconda.sh -b -p "$HOME/miniconda3" >"$LOG_PATH" 2>&1
|
||||
|
||||
if [ $? -eq 0 ]; then
|
||||
echo "== Miniconda Installed =="
|
||||
else
|
||||
echo "Failed to Install miniconda"
|
||||
tail -n 50 "$LOG_PATH"
|
||||
exit 1
|
||||
fi
|
||||
|
||||
rm miniconda.sh
|
||||
|
||||
source "$HOME/miniconda3/etc/profile.d/conda.sh"
|
||||
|
||||
"$HOME/miniconda3/bin/conda" config --add channels conda-forge
|
||||
|
||||
"$HOME/miniconda3/bin/conda" update -q --all -y 1>/dev/null
|
||||
|
||||
"$HOME/miniconda3/bin/conda" install python=3.11 -q -y
|
||||
|
||||
"$HOME/miniconda3/bin/conda" install gcc=14 gxx ffmpeg cmake make unzip -q -y
|
||||
|
||||
if [ "$CUDA_VERSION" = "12.8" ]; then
|
||||
"$HOME/miniconda3/bin/pip" install torch torchaudio --no-cache-dir --index-url https://download.pytorch.org/whl/cu128
|
||||
elif [ "$CUDA_VERSION" = "12.6" ]; then
|
||||
"$HOME/miniconda3/bin/pip" install torch==2.6 torchaudio --no-cache-dir --index-url https://download.pytorch.org/whl/cu126
|
||||
fi
|
||||
|
||||
"$HOME/miniconda3/bin/pip" cache purge
|
||||
|
||||
rm $LOG_PATH
|
||||
|
||||
rm -rf "$HOME/miniconda3/pkgs"
|
||||
|
||||
mkdir -p "$HOME/miniconda3/pkgs"
|
||||
|
||||
rm -rf "$HOME/.conda" "$HOME/.cache"
|
||||
@@ -0,0 +1,82 @@
|
||||
#!/bin/bash
|
||||
|
||||
set -e
|
||||
|
||||
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" &>/dev/null && pwd)"
|
||||
|
||||
cd "$SCRIPT_DIR" || exit 1
|
||||
|
||||
cd .. || exit 1
|
||||
|
||||
if [ -d "$HOME/conda" ]; then
|
||||
exit 0
|
||||
fi
|
||||
|
||||
WORKFLOW=${WORKFLOW:-"false"}
|
||||
TARGETPLATFORM=${TARGETPLATFORM:-"linux/amd64"}
|
||||
|
||||
if [ "$WORKFLOW" = "true" ]; then
|
||||
WGET_CMD=(wget -nv --tries=25 --wait=5 --read-timeout=40 --retry-on-http-error=404)
|
||||
else
|
||||
WGET_CMD=(wget --tries=25 --wait=5 --read-timeout=40 --retry-on-http-error=404)
|
||||
fi
|
||||
|
||||
if [ "$TARGETPLATFORM" = "linux/amd64" ]; then
|
||||
"${WGET_CMD[@]}" -O Miniforge.sh "https://github.com/conda-forge/miniforge/releases/latest/download/Miniforge3-$(uname)-$(uname -m).sh"
|
||||
SYSROOT_PKG="sysroot_linux-64>=2.28"
|
||||
elif [ "$TARGETPLATFORM" = "linux/arm64" ]; then
|
||||
"${WGET_CMD[@]}" -O Miniforge.sh "https://github.com/conda-forge/miniforge/releases/latest/download/Miniforge3-$(uname)-$(uname -m).sh"
|
||||
SYSROOT_PKG="sysroot_linux-aarch64>=2.28"
|
||||
else
|
||||
exit 1
|
||||
fi
|
||||
|
||||
LOG_PATH="/tmp/miniforge-install.log"
|
||||
|
||||
bash Miniforge.sh -b -p "$HOME/conda" >"$LOG_PATH" 2>&1
|
||||
|
||||
if [ $? -eq 0 ]; then
|
||||
echo "== Miniforge Installed =="
|
||||
else
|
||||
echo "Failed to Install miniforge"
|
||||
tail -n 50 "$LOG_PATH"
|
||||
exit 1
|
||||
fi
|
||||
|
||||
rm Miniforge.sh
|
||||
|
||||
source "$HOME/conda/etc/profile.d/conda.sh"
|
||||
|
||||
"$HOME/conda/bin/conda" init bash
|
||||
|
||||
source "$HOME/.bashrc"
|
||||
|
||||
"$HOME/conda/bin/conda" info
|
||||
|
||||
"$HOME/conda/bin/conda" update --all -y
|
||||
|
||||
"$HOME/conda/bin/conda" install python=3.12 -y
|
||||
|
||||
"$HOME/conda/bin/conda" install gcc=11 gxx ffmpeg cmake make unzip $SYSROOT_PKG "libstdcxx-ng>=11" -y
|
||||
|
||||
if [ "$CUDA_VERSION" = "12.8" ]; then
|
||||
"$HOME/conda/bin/pip" install torch torchcodec --no-cache-dir --index-url https://download.pytorch.org/whl/cu128
|
||||
"$HOME/conda/bin/conda" install cuda-nvcc=12.8 -y
|
||||
elif [ "$CUDA_VERSION" = "12.6" ]; then
|
||||
"$HOME/conda/bin/pip" install torch torchcodec --no-cache-dir --index-url https://download.pytorch.org/whl/cu126
|
||||
"$HOME/conda/bin/conda" install cuda-nvcc=12.6 -y
|
||||
fi
|
||||
|
||||
export PATH="$HOME/conda/bin:$PATH"
|
||||
|
||||
"$HOME/conda/bin/pip" install psutil ninja packaging wheel "setuptools>=42" einops
|
||||
"$HOME/conda/bin/pip" install flash-attn -i https://xxxxrt666.github.io/PIP-Index/ --no-build-isolation
|
||||
"$HOME/conda/bin/pip" cache purge
|
||||
|
||||
rm $LOG_PATH
|
||||
|
||||
rm -rf "$HOME/conda/pkgs"
|
||||
|
||||
mkdir -p "$HOME/conda/pkgs"
|
||||
|
||||
rm -rf "$HOME/.conda" "$HOME/.cache"
|
||||
+1
-3
@@ -4,7 +4,7 @@ ARG TORCH_BASE=full
|
||||
FROM xxxxrt666/torch-base:cu${CUDA_VERSION}-${TORCH_BASE}
|
||||
|
||||
LABEL maintainer="XXXXRT"
|
||||
LABEL version="V4"
|
||||
LABEL version="V2 Pro"
|
||||
LABEL description="Docker image for GPT-SoVITS"
|
||||
|
||||
ARG CUDA_VERSION=12.6
|
||||
@@ -26,8 +26,6 @@ ENV WORKFLOW=${WORKFLOW}
|
||||
ARG TARGETPLATFORM
|
||||
ENV TARGETPLATFORM=${TARGETPLATFORM}
|
||||
|
||||
RUN bash Docker/miniconda_install.sh
|
||||
|
||||
COPY extra-req.txt /workspace/GPT-SoVITS/
|
||||
|
||||
COPY requirements.txt /workspace/GPT-SoVITS/
|
||||
|
||||
@@ -67,8 +67,10 @@ class Text2SemanticDataset(Dataset):
|
||||
)
|
||||
) # "%s/3-bert"%exp_dir#bert_dir
|
||||
self.path6 = semantic_path # "%s/6-name2semantic.tsv"%exp_dir#semantic_path
|
||||
assert os.path.exists(self.path2)
|
||||
assert os.path.exists(self.path6)
|
||||
if not os.path.exists(self.path2):
|
||||
raise FileNotFoundError(f"Phoneme data file not found: {self.path2}")
|
||||
if not os.path.exists(self.path6):
|
||||
raise FileNotFoundError(f"Semantic data file not found: {self.path6}")
|
||||
self.phoneme_data = {}
|
||||
with open(self.path2, "r", encoding="utf8") as f:
|
||||
lines = f.read().strip("\n").split("\n")
|
||||
@@ -131,7 +133,7 @@ class Text2SemanticDataset(Dataset):
|
||||
phoneme, word2ph, text = self.phoneme_data[item_name]
|
||||
except Exception:
|
||||
traceback.print_exc()
|
||||
# print(f"{item_name} not in self.phoneme_data !")
|
||||
print(f"Warning: File \"{item_name}\" not in self.phoneme_data! Skipped. ")
|
||||
num_not_in += 1
|
||||
continue
|
||||
|
||||
@@ -152,7 +154,7 @@ class Text2SemanticDataset(Dataset):
|
||||
phoneme_ids = cleaned_text_to_sequence(phoneme, version)
|
||||
except:
|
||||
traceback.print_exc()
|
||||
# print(f"{item_name} not in self.phoneme_data !")
|
||||
print(f"Warning: Failed to convert phonemes to sequence for file \"{item_name}\"! Skipped. ")
|
||||
num_not_in += 1
|
||||
continue
|
||||
# if len(phoneme_ids) >400:###########2:改为恒定限制为semantic/2.5就行
|
||||
@@ -228,7 +230,11 @@ class Text2SemanticDataset(Dataset):
|
||||
# bert_feature=torch.zeros_like(phoneme_ids,dtype=torch.float32)
|
||||
bert_feature = None
|
||||
else:
|
||||
assert bert_feature.shape[-1] == len(phoneme_ids)
|
||||
try:
|
||||
assert bert_feature.shape[-1] == len(phoneme_ids)
|
||||
except AssertionError:
|
||||
print(f"AssertionError: The BERT feature dimension ({bert_feature.shape[-1]}) of the file '{item_name}' does not match the length of the phoneme sequence ({len(phoneme_ids)}).")
|
||||
raise
|
||||
return {
|
||||
"idx": idx,
|
||||
"phoneme_ids": phoneme_ids,
|
||||
|
||||
@@ -0,0 +1,76 @@
|
||||
import math
|
||||
|
||||
import torch
|
||||
from torch import nn
|
||||
|
||||
|
||||
class TokenEmbedding(nn.Module):
|
||||
def __init__(self, embedding_dim: int, vocab_size: int, dropout: float = 0.0):
|
||||
super().__init__()
|
||||
self.vocab_size = vocab_size
|
||||
self.embedding_dim = embedding_dim
|
||||
self.dropout = nn.Dropout(p=dropout)
|
||||
self.word_embeddings = nn.Embedding(self.vocab_size, self.embedding_dim)
|
||||
|
||||
@property
|
||||
def weight(self) -> torch.Tensor:
|
||||
return self.word_embeddings.weight
|
||||
|
||||
def embedding(self, index: int) -> torch.Tensor:
|
||||
return self.word_embeddings.weight[index : index + 1]
|
||||
|
||||
def forward(self, x: torch.Tensor):
|
||||
x = self.word_embeddings(x)
|
||||
x = self.dropout(x)
|
||||
return x
|
||||
|
||||
|
||||
class SinePositionalEmbeddingNested(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
embedding_dim: int,
|
||||
dropout: float = 0.0,
|
||||
scale: bool = False,
|
||||
alpha: bool = False,
|
||||
max_batch_size: int = 20,
|
||||
max_seq_len: int = 2500,
|
||||
):
|
||||
super().__init__()
|
||||
self.embedding_dim = embedding_dim
|
||||
self.x_scale = math.sqrt(embedding_dim) if scale else 1.0
|
||||
self.alpha = nn.Parameter(torch.ones(1), requires_grad=alpha)
|
||||
self.dropout = nn.Dropout(p=dropout)
|
||||
self.max_batch_size = max_batch_size
|
||||
self.max_seq_len = max_seq_len
|
||||
|
||||
self.reverse = False
|
||||
self.register_buffer(
|
||||
"pe", torch.zeros(max_batch_size, max_seq_len, embedding_dim), persistent=False
|
||||
)
|
||||
self.pe: torch.Tensor
|
||||
self.compute_pe()
|
||||
|
||||
def compute_pe(self):
|
||||
if self.reverse:
|
||||
position = torch.arange(self.max_seq_len - 1, -1, -1.0, dtype=torch.float32).unsqueeze(1)
|
||||
else:
|
||||
position = torch.arange(self.max_seq_len, dtype=torch.float32).unsqueeze(1)
|
||||
div_term = torch.exp(
|
||||
torch.arange(0, self.embedding_dim, 2, dtype=torch.float32)
|
||||
* -(math.log(10000.0) / self.embedding_dim)
|
||||
)
|
||||
pe = self.pe
|
||||
pe[:, :, 0::2] = torch.sin(position * div_term)
|
||||
pe[:, :, 1::2] = torch.cos(position * div_term)
|
||||
|
||||
def forward(self, input_pos: torch.Tensor, x: torch.Tensor) -> torch.Tensor:
|
||||
batch_size = x.shape[0]
|
||||
pe_values = self.pe[torch.arange(batch_size), input_pos - 1]
|
||||
return x * self.x_scale + self.alpha * pe_values.unsqueeze(1)
|
||||
|
||||
def prefill(self, x: torch.Tensor) -> torch.Tensor:
|
||||
input_pos = torch.tensor([i.shape[0] for i in x.unbind()])
|
||||
pe_values = torch.nested.nested_tensor(
|
||||
[self.pe[i, : input_pos[i], :] for i in range(input_pos.size(0))]
|
||||
)
|
||||
return x * self.x_scale + self.alpha.item() * pe_values
|
||||
@@ -0,0 +1,78 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass
|
||||
from typing import List, Literal, Optional
|
||||
|
||||
import torch
|
||||
|
||||
Tensor = torch.Tensor
|
||||
|
||||
|
||||
@dataclass
|
||||
class T2SResult:
|
||||
result: List[Tensor] | None = None
|
||||
infer_speed: float = 0.0
|
||||
status: Literal["Success", "Error"] = "Success"
|
||||
exception: Optional[Exception] = None
|
||||
traceback: Optional[str] = None
|
||||
|
||||
|
||||
@dataclass
|
||||
class T2SRequest:
|
||||
x: List[torch.Tensor]
|
||||
x_lens: Tensor
|
||||
prompts: torch.Tensor
|
||||
bert_feature: List[Tensor]
|
||||
valid_length: int
|
||||
top_k: int = 5
|
||||
top_p: float = 1
|
||||
early_stop_num: int = -1
|
||||
temperature: float = 1.0
|
||||
repetition_penalty: float = 1.35
|
||||
use_cuda_graph: bool = False
|
||||
debug: bool = False
|
||||
|
||||
|
||||
class T2SSession:
|
||||
def __init__(self, decoder, request: T2SRequest, device: torch.device, dtype: torch.dtype):
|
||||
with device:
|
||||
self.decoder = decoder
|
||||
self.request = request
|
||||
self.device = device
|
||||
self.dtype = dtype
|
||||
|
||||
bsz = len(request.x)
|
||||
y_len = request.prompts.size(-1)
|
||||
self.bsz = bsz
|
||||
self.y_len = y_len
|
||||
|
||||
from AR.models.t2s_model_cudagraph import Sampler
|
||||
|
||||
self.sampler = Sampler(bsz, decoder.vocab_size)
|
||||
|
||||
self.x = request.x
|
||||
self.x_lens = request.x_lens.to(torch.int32)
|
||||
self.y = request.prompts
|
||||
self.bert_feature = request.bert_feature
|
||||
|
||||
self.prefill_len = self.x_lens + self.y.size(1)
|
||||
|
||||
self.input_pos = torch.zeros_like(self.prefill_len)
|
||||
self.input_pos.add_(self.prefill_len)
|
||||
|
||||
self.completed = torch.Tensor([False] * len(self.x)).bool().to(device)
|
||||
self.y_results: List[Tensor] = [None] * len(self.x) # type: ignore
|
||||
|
||||
self.xy_pos = decoder.embed(self.x, self.y, self.bert_feature)
|
||||
|
||||
attn_mask = []
|
||||
for bs in range(bsz):
|
||||
pos = int(self.x_lens[bs].item())
|
||||
mask = torch.zeros(pos + y_len, pos + y_len).bool()
|
||||
mask[:, :pos].fill_(True)
|
||||
if y_len > 0:
|
||||
mask[-y_len:, -y_len:] = ~torch.triu(
|
||||
torch.ones(y_len, y_len, dtype=torch.bool), diagonal=1
|
||||
)
|
||||
attn_mask.append(mask)
|
||||
self.attn_mask_nested = torch.nested.nested_tensor(attn_mask)
|
||||
@@ -28,7 +28,8 @@ class Text2SemanticLightningModule(LightningModule):
|
||||
self.load_state_dict(
|
||||
torch.load(
|
||||
pretrained_s1,
|
||||
map_location="cpu", weights_only=False,
|
||||
map_location="cpu",
|
||||
weights_only=False,
|
||||
)["weight"],
|
||||
)
|
||||
)
|
||||
|
||||
@@ -356,7 +356,7 @@ class Text2SemanticDecoder(nn.Module):
|
||||
x = self.ar_text_embedding(x)
|
||||
x = x + self.bert_proj(bert_feature.transpose(1, 2))
|
||||
x = self.ar_text_position(x)
|
||||
x_mask = make_pad_mask(x_lens)
|
||||
x_mask = make_pad_mask_left(x_lens)
|
||||
|
||||
y_mask = make_pad_mask(y_lens)
|
||||
y_mask_int = y_mask.type(torch.int64)
|
||||
@@ -420,7 +420,7 @@ class Text2SemanticDecoder(nn.Module):
|
||||
mask=xy_attn_mask,
|
||||
)
|
||||
x_len = x_lens.max()
|
||||
logits = self.ar_predict_layer(xy_dec[:, x_len:])
|
||||
logits = self.ar_predict_layer(xy_dec[:, x_len-1:])
|
||||
|
||||
###### DPO #############
|
||||
reject_xy_pos, reject_xy_attn_mask, reject_targets = self.make_input_data(
|
||||
@@ -432,7 +432,7 @@ class Text2SemanticDecoder(nn.Module):
|
||||
mask=reject_xy_attn_mask,
|
||||
)
|
||||
x_len = x_lens.max()
|
||||
reject_logits = self.ar_predict_layer(reject_xy_dec[:, x_len:])
|
||||
reject_logits = self.ar_predict_layer(reject_xy_dec[:, x_len-1:])
|
||||
|
||||
# loss
|
||||
# from feiteng: 每次 duration 越多, 梯度更新也应该更多, 所以用 sum
|
||||
@@ -455,7 +455,7 @@ class Text2SemanticDecoder(nn.Module):
|
||||
x = self.ar_text_embedding(x)
|
||||
x = x + self.bert_proj(bert_feature.transpose(1, 2))
|
||||
x = self.ar_text_position(x)
|
||||
x_mask = make_pad_mask(x_lens)
|
||||
x_mask = make_pad_mask_left(x_lens)
|
||||
|
||||
y_mask = make_pad_mask(y_lens)
|
||||
y_mask_int = y_mask.type(torch.int64)
|
||||
@@ -502,7 +502,7 @@ class Text2SemanticDecoder(nn.Module):
|
||||
(xy_pos, None),
|
||||
mask=xy_attn_mask,
|
||||
)
|
||||
logits = self.ar_predict_layer(xy_dec[:, x_len:]).permute(0, 2, 1)
|
||||
logits = self.ar_predict_layer(xy_dec[:, x_len-1:]).permute(0, 2, 1)
|
||||
# loss
|
||||
# from feiteng: 每次 duration 越多, 梯度更新也应该更多, 所以用 sum
|
||||
loss = F.cross_entropy(logits, targets, reduction="sum")
|
||||
@@ -578,7 +578,7 @@ class Text2SemanticDecoder(nn.Module):
|
||||
def pad_y_eos(self, y, y_mask_int, eos_id):
|
||||
targets = F.pad(y, (0, 1), value=0) + eos_id * F.pad(y_mask_int, (0, 1), value=1)
|
||||
# 错位
|
||||
return targets[:, :-1], targets[:, 1:]
|
||||
return targets[:, :-1], targets
|
||||
|
||||
def infer_panel_batch_infer(
|
||||
self,
|
||||
@@ -707,10 +707,12 @@ class Text2SemanticDecoder(nn.Module):
|
||||
|
||||
if idx == 0:
|
||||
attn_mask = F.pad(attn_mask[:, :, -1].unsqueeze(-2), (0, 1), value=False)
|
||||
logits = logits[:, :-1]
|
||||
else:
|
||||
attn_mask = F.pad(attn_mask, (0, 1), value=False)
|
||||
|
||||
if idx < 11: ###至少预测出10个token不然不给停止(0.4s)
|
||||
logits = logits[:, :-1]
|
||||
|
||||
samples = sample(
|
||||
logits, y, top_k=top_k, top_p=top_p, repetition_penalty=repetition_penalty, temperature=temperature
|
||||
)[0]
|
||||
@@ -794,7 +796,7 @@ class Text2SemanticDecoder(nn.Module):
|
||||
y_list = []
|
||||
idx_list = []
|
||||
for i in range(len(x)):
|
||||
y, idx = self.infer_panel_naive(
|
||||
y, idx = next(self.infer_panel_naive(
|
||||
x[i].unsqueeze(0),
|
||||
x_lens[i],
|
||||
prompts[i].unsqueeze(0) if prompts is not None else None,
|
||||
@@ -805,7 +807,7 @@ class Text2SemanticDecoder(nn.Module):
|
||||
temperature,
|
||||
repetition_penalty,
|
||||
**kwargs,
|
||||
)
|
||||
))
|
||||
y_list.append(y[0])
|
||||
idx_list.append(idx)
|
||||
|
||||
@@ -822,8 +824,15 @@ class Text2SemanticDecoder(nn.Module):
|
||||
early_stop_num: int = -1,
|
||||
temperature: float = 1.0,
|
||||
repetition_penalty: float = 1.35,
|
||||
streaming_mode: bool = False,
|
||||
chunk_length: int = 24,
|
||||
**kwargs,
|
||||
):
|
||||
mute_emb_sim_matrix = kwargs.get("mute_emb_sim_matrix", None)
|
||||
chunk_split_thershold = kwargs.get("chunk_split_thershold", 0.3)
|
||||
check_token_num = 2
|
||||
|
||||
|
||||
x = self.ar_text_embedding(x)
|
||||
x = x + self.bert_proj(bert_feature.transpose(1, 2))
|
||||
x = self.ar_text_position(x)
|
||||
@@ -875,7 +884,10 @@ class Text2SemanticDecoder(nn.Module):
|
||||
.to(device=x.device, dtype=torch.bool)
|
||||
)
|
||||
|
||||
token_counter = 0
|
||||
curr_ptr = prefix_len
|
||||
for idx in tqdm(range(1500)):
|
||||
token_counter+=1
|
||||
if xy_attn_mask is not None:
|
||||
xy_dec, k_cache, v_cache = self.t2s_transformer.process_prompt(xy_pos, xy_attn_mask, None)
|
||||
else:
|
||||
@@ -900,22 +912,56 @@ class Text2SemanticDecoder(nn.Module):
|
||||
|
||||
if torch.argmax(logits, dim=-1)[0] == self.EOS or samples[0, 0] == self.EOS:
|
||||
stop = True
|
||||
y=y[:, :-1]
|
||||
token_counter -= 1
|
||||
|
||||
if idx == 1499:
|
||||
stop = True
|
||||
|
||||
if stop:
|
||||
if y.shape[1] == 0:
|
||||
y = torch.concat([y, torch.zeros_like(samples)], dim=1)
|
||||
print("bad zero prediction")
|
||||
print(f"T2S Decoding EOS [{prefix_len} -> {y.shape[1]}]")
|
||||
# print(f"T2S Decoding EOS [{prefix_len} -> {y.shape[1]}]")
|
||||
if streaming_mode:
|
||||
yield y[:, curr_ptr:] if curr_ptr<y.shape[1] else None, True
|
||||
break
|
||||
|
||||
|
||||
if streaming_mode and (mute_emb_sim_matrix is not None) and (token_counter >= chunk_length+check_token_num):
|
||||
score = mute_emb_sim_matrix[y[0, curr_ptr:]] - chunk_split_thershold
|
||||
score[score<0]=-1
|
||||
score[:-1]=score[:-1]+score[1:] ##考虑连续两个token
|
||||
argmax_idx = score.argmax()
|
||||
|
||||
if score[argmax_idx]>=0 and argmax_idx+1>=chunk_length:
|
||||
print(f"\n\ncurr_ptr:{curr_ptr}")
|
||||
yield y[:, curr_ptr:], False
|
||||
token_counter -= argmax_idx+1
|
||||
curr_ptr += argmax_idx+1
|
||||
|
||||
|
||||
elif streaming_mode and (mute_emb_sim_matrix is None) and (token_counter >= chunk_length):
|
||||
yield y[:, -token_counter:], False
|
||||
curr_ptr+=token_counter
|
||||
token_counter = 0
|
||||
|
||||
|
||||
|
||||
####################### update next step ###################################
|
||||
y_emb = self.ar_audio_embedding(y[:, -1:])
|
||||
xy_pos = y_emb * self.ar_audio_position.x_scale + self.ar_audio_position.alpha * self.ar_audio_position.pe[
|
||||
:, y_len + idx
|
||||
].to(dtype=y_emb.dtype, device=y_emb.device)
|
||||
|
||||
if ref_free:
|
||||
return y[:, :-1], 0
|
||||
return y[:, :-1], idx
|
||||
|
||||
|
||||
if not streaming_mode:
|
||||
if ref_free:
|
||||
yield y, 0
|
||||
yield y, idx
|
||||
|
||||
|
||||
|
||||
def infer_panel(
|
||||
self,
|
||||
@@ -930,6 +976,6 @@ class Text2SemanticDecoder(nn.Module):
|
||||
repetition_penalty: float = 1.35,
|
||||
**kwargs,
|
||||
):
|
||||
return self.infer_panel_naive(
|
||||
return next(self.infer_panel_naive(
|
||||
x, x_lens, prompts, bert_feature, top_k, top_p, early_stop_num, temperature, repetition_penalty, **kwargs
|
||||
)
|
||||
))
|
||||
|
||||
@@ -0,0 +1,602 @@
|
||||
"""
|
||||
CUDA Graph accelerated T2S decoder.
|
||||
Uses PyTorch native scaled_dot_product_attention (no flash_attn dependency).
|
||||
Adapted from gsvpp/AR/models/t2s_model_abc.py and t2s_model_flash_attn.py.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
import time
|
||||
import traceback
|
||||
from typing import Dict, List, MutableSequence, Optional, Tuple
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
from torch.cuda.graphs import CUDAGraph
|
||||
from tqdm import tqdm
|
||||
|
||||
from AR.models.embedding_cudagraph import (
|
||||
SinePositionalEmbeddingNested as SinePositionalEmbedding,
|
||||
)
|
||||
from AR.models.embedding_cudagraph import TokenEmbedding
|
||||
from AR.models.structs_cudagraph import T2SRequest, T2SResult, T2SSession
|
||||
|
||||
Tensor = torch.Tensor
|
||||
|
||||
|
||||
class Sampler(nn.Module):
|
||||
def __init__(self, batch_size: int, vocab_size: int) -> None:
|
||||
super().__init__()
|
||||
self.batch_size = batch_size
|
||||
|
||||
def sample(
|
||||
self,
|
||||
logits: Tensor,
|
||||
previous_tokens: Tensor,
|
||||
temperature: float,
|
||||
top_k: int,
|
||||
top_p: float,
|
||||
repetition_penalty: float,
|
||||
) -> Tensor:
|
||||
previous_tokens = previous_tokens.long()
|
||||
score = torch.gather(logits, dim=1, index=previous_tokens)
|
||||
score = torch.where(
|
||||
score < 0, score * repetition_penalty, score / repetition_penalty
|
||||
)
|
||||
logits.scatter_(dim=1, index=previous_tokens, src=score)
|
||||
|
||||
sorted_logits, sorted_indices = torch.sort(logits, descending=True)
|
||||
cum_probs = torch.cumsum(
|
||||
torch.nn.functional.softmax(sorted_logits, dim=-1), dim=-1
|
||||
)
|
||||
sorted_indices_to_remove = cum_probs > top_p
|
||||
sorted_indices_to_remove[:, 0] = False
|
||||
indices_to_remove = sorted_indices_to_remove.scatter(
|
||||
dim=1, index=sorted_indices, src=sorted_indices_to_remove
|
||||
)
|
||||
logits = logits.masked_fill(indices_to_remove, -float("Inf"))
|
||||
|
||||
logits = logits / max(temperature, 1e-5)
|
||||
|
||||
v, _ = torch.topk(logits, top_k)
|
||||
pivot = v[:, -1].unsqueeze(-1)
|
||||
logits = torch.where(logits < pivot, -float("Inf"), logits)
|
||||
|
||||
probs = torch.nn.functional.softmax(logits, dim=-1)
|
||||
q = torch.empty_like(probs).exponential_(1.0)
|
||||
idx_next = torch.argmax(probs / q, dim=-1, keepdim=True).to(dtype=torch.int32)
|
||||
|
||||
return idx_next
|
||||
|
||||
|
||||
# ─── KV Cache ────────────────────���───────────────────────────────────────────
|
||||
|
||||
|
||||
class KVCacheNHD(nn.Module):
|
||||
def __init__(self, batch_size, max_seq_length, n_heads, head_dim):
|
||||
super().__init__()
|
||||
assert batch_size > 0
|
||||
cache_shape = (batch_size, max_seq_length, n_heads, head_dim)
|
||||
self.n_head = n_heads
|
||||
self.head_dim = head_dim
|
||||
self.batch_size = batch_size
|
||||
self.max_seq_length = max_seq_length
|
||||
self.register_buffer(
|
||||
"k_cache", torch.zeros(size=cache_shape), persistent=False
|
||||
)
|
||||
self.register_buffer(
|
||||
"v_cache", torch.zeros(size=cache_shape), persistent=False
|
||||
)
|
||||
|
||||
def update(self, input_pos: Tensor, k_val: Tensor, v_val: Tensor):
|
||||
index = (
|
||||
(input_pos - 1)
|
||||
.unsqueeze(-1)
|
||||
.unsqueeze(-1)
|
||||
.unsqueeze(-1)
|
||||
.expand(-1, -1, self.n_head, self.head_dim)
|
||||
.to(torch.int64)
|
||||
)
|
||||
k_out = self.k_cache
|
||||
v_out = self.v_cache
|
||||
k_out.scatter_(1, index, k_val)
|
||||
v_out.scatter_(1, index, v_val)
|
||||
return k_out, v_out
|
||||
|
||||
def empty(self):
|
||||
self.k_cache.zero_()
|
||||
self.v_cache.zero_()
|
||||
|
||||
def prefill_kv(self, k_val: Tensor, v_val: Tensor, bs: int):
|
||||
self.k_cache[[bs], : k_val.shape[1]] = k_val
|
||||
self.v_cache[[bs], : v_val.shape[1]] = v_val
|
||||
|
||||
|
||||
# ─── Attention (PyTorch native SDPA, no flash_attn) ─────────────────────────
|
||||
|
||||
|
||||
class Attention(nn.Module):
|
||||
def __init__(self, n_head: int, hidden_dim: int):
|
||||
super().__init__()
|
||||
self.n_head = n_head
|
||||
self.hidden_dim = hidden_dim
|
||||
assert hidden_dim % n_head == 0
|
||||
self.head_dim = hidden_dim // n_head
|
||||
self.in_proj = nn.Linear(hidden_dim, hidden_dim * 3, bias=True)
|
||||
self.out_proj = nn.Linear(hidden_dim, hidden_dim, bias=True)
|
||||
self.dropout = nn.Dropout(0.1)
|
||||
|
||||
self._register_load_state_dict_pre_hook(self.load_hook)
|
||||
|
||||
def load_hook(self, state_dict: dict, prefix, *args):
|
||||
keys_to_modify = [key for key in state_dict if "in_proj_" in key]
|
||||
for key in keys_to_modify:
|
||||
new_key = key.replace("in_proj_", "in_proj.")
|
||||
state_dict[new_key] = state_dict.pop(key)
|
||||
|
||||
def forward(
|
||||
self, x: Tensor, input_pos: Tensor, kv_cache: KVCacheNHD
|
||||
) -> Tensor:
|
||||
bsz, seqlen, _ = x.shape
|
||||
|
||||
q, k, v = self.in_proj.forward(x).chunk(3, dim=-1)
|
||||
|
||||
q = q.view(bsz, seqlen, self.n_head, self.head_dim)
|
||||
k = k.view(bsz, seqlen, self.n_head, self.head_dim)
|
||||
v = v.view(bsz, seqlen, self.n_head, self.head_dim)
|
||||
|
||||
k_cache, v_cache = kv_cache.update(input_pos, k, v)
|
||||
|
||||
q = q.transpose(1, 2) # [B, H, 1, D]
|
||||
k_out = k_cache.transpose(1, 2) # [B, H, max_seq, D]
|
||||
v_out = v_cache.transpose(1, 2) # [B, H, max_seq, D]
|
||||
|
||||
attn = F.scaled_dot_product_attention(q, k_out, v_out)
|
||||
|
||||
attn = self.dropout.forward(attn)
|
||||
attn = attn.transpose(1, 2).reshape(bsz, seqlen, self.hidden_dim)
|
||||
attn = self.out_proj.forward(attn)
|
||||
return attn
|
||||
|
||||
def prefill(self, x: Tensor, mask: Tensor, kv_cache: KVCacheNHD) -> Tensor:
|
||||
bsz = x.size(0)
|
||||
outputs = []
|
||||
for bs in range(bsz):
|
||||
x_b = x[bs].unsqueeze(0)
|
||||
q, k, v = self.in_proj.forward(x_b.unsqueeze(0)).chunk(3, dim=-1)
|
||||
q = q.contiguous().view(1, -1, self.n_head, self.head_dim)
|
||||
k = k.contiguous().view(1, -1, self.n_head, self.head_dim)
|
||||
v = v.contiguous().view(1, -1, self.n_head, self.head_dim)
|
||||
kv_cache.prefill_kv(k, v, bs)
|
||||
q, k, v = map(lambda t: t.transpose(1, 2), (q, k, v))
|
||||
attn_mask = (
|
||||
mask[bs].unsqueeze(0).unsqueeze(0).expand(1, self.n_head, -1, -1)
|
||||
)
|
||||
attn = F.scaled_dot_product_attention(q, k, v, attn_mask=attn_mask)
|
||||
attn = self.dropout.forward(attn)
|
||||
attn = attn.transpose(1, 2).contiguous().view(1, -1, self.hidden_dim)
|
||||
output = self.out_proj.forward(attn)
|
||||
outputs.append(output.squeeze(0))
|
||||
return torch.nested.nested_tensor(outputs)
|
||||
|
||||
|
||||
# ─── Feed Forward ────────────────────────────────────────────────────────────
|
||||
|
||||
|
||||
class FeedForward(nn.Module):
|
||||
def __init__(self, dim: int, hidden_dim: int) -> None:
|
||||
super().__init__()
|
||||
self.linear1 = nn.Linear(dim, hidden_dim, bias=True)
|
||||
self.linear2 = nn.Linear(hidden_dim, dim, bias=True)
|
||||
self.dropout = nn.Dropout(0.1)
|
||||
|
||||
def forward(self, x: Tensor) -> Tensor:
|
||||
return self.dropout.forward(
|
||||
self.linear2(self.dropout.forward(F.relu(self.linear1(x))))
|
||||
)
|
||||
|
||||
|
||||
# ─── Transformer Block ──────────────────────────────────────────────────────
|
||||
|
||||
|
||||
class TransformerBlock(nn.Module):
|
||||
def __init__(self, n_head, ffn_dim, hidden_dim) -> None:
|
||||
super().__init__()
|
||||
self.hidden_dim = hidden_dim
|
||||
self.attention = Attention(n_head, hidden_dim)
|
||||
self.feed_forward = FeedForward(hidden_dim, ffn_dim)
|
||||
self.attention_norm = nn.LayerNorm([hidden_dim])
|
||||
self.ffn_norm = nn.LayerNorm([hidden_dim])
|
||||
self.dropout = nn.Dropout(0.1)
|
||||
|
||||
self._register_load_state_dict_pre_hook(self.load_hook)
|
||||
|
||||
def load_hook(self, state_dict: dict[str, Tensor], prefix, *args):
|
||||
for key in list(state_dict.keys()):
|
||||
new_key = (
|
||||
key.replace("self_attn", "attention")
|
||||
.replace("linear", "feed_forward.linear")
|
||||
.replace("norm1", "attention_norm")
|
||||
.replace("norm2", "ffn_norm")
|
||||
)
|
||||
state_dict[new_key] = state_dict.pop(key)
|
||||
|
||||
def forward(
|
||||
self, x: Tensor, input_pos: Tensor, kv_cache: KVCacheNHD
|
||||
) -> Tensor:
|
||||
h = self.attention_norm.forward(
|
||||
x + self.dropout.forward(self.attention.forward(x, input_pos, kv_cache))
|
||||
)
|
||||
out = self.ffn_norm.forward(h + self.feed_forward.forward(h))
|
||||
return out
|
||||
|
||||
def prefill(self, x: Tensor, mask: Tensor, kv_cache: KVCacheNHD) -> Tensor:
|
||||
h = self.attention_norm.forward(
|
||||
x + self.dropout.forward(self.attention.prefill(x, mask, kv_cache))
|
||||
)
|
||||
out = self.ffn_norm.forward(h + self.feed_forward.forward(h))
|
||||
return out
|
||||
|
||||
|
||||
# ─── Transformer Decoder ────────────────────────────────────────────────────
|
||||
|
||||
|
||||
class TransformerDecoder(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
hidden_dim,
|
||||
n_layer,
|
||||
n_head,
|
||||
ffn_dim,
|
||||
vocab_size,
|
||||
max_seq_length,
|
||||
max_batch_size,
|
||||
) -> None:
|
||||
super().__init__()
|
||||
self.hidden_dim = hidden_dim
|
||||
self.n_head = n_head
|
||||
assert hidden_dim % n_head == 0
|
||||
self.head_dim = hidden_dim // n_head
|
||||
self.vocab_size = vocab_size
|
||||
self.n_layer = n_layer
|
||||
self.layers = nn.ModuleList(
|
||||
TransformerBlock(n_head, ffn_dim, hidden_dim) for _ in range(n_layer)
|
||||
)
|
||||
self.max_seq_length: int = max_seq_length
|
||||
self.max_batch_size: int = max_batch_size
|
||||
|
||||
def forward(
|
||||
self,
|
||||
input_pos: Tensor,
|
||||
x: Tensor,
|
||||
kv_caches: MutableSequence[KVCacheNHD],
|
||||
):
|
||||
for layer, kv_cache in zip(self.layers, kv_caches):
|
||||
x = layer.forward(x, input_pos, kv_cache)
|
||||
return x
|
||||
|
||||
def prefill(
|
||||
self,
|
||||
x: Tensor,
|
||||
mask: Tensor,
|
||||
kv_caches: MutableSequence[KVCacheNHD],
|
||||
):
|
||||
for layer, kv_cache in zip(self.layers, kv_caches):
|
||||
x = layer.prefill(x, mask, kv_cache)
|
||||
return x
|
||||
|
||||
|
||||
# ─── T2S Decoder ─────────────────────────────────────────────────────────────
|
||||
|
||||
|
||||
class T2SDecoder(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
config,
|
||||
*args,
|
||||
norm_first=False,
|
||||
max_seq_length=2500,
|
||||
max_batch_size=10,
|
||||
**kwds,
|
||||
) -> None:
|
||||
super().__init__()
|
||||
hidden_dim = config["model"]["hidden_dim"]
|
||||
embedding_dim = config["model"]["embedding_dim"]
|
||||
n_head = config["model"]["head"]
|
||||
n_layer = config["model"]["n_layer"]
|
||||
vocab_size = config["model"]["vocab_size"]
|
||||
phoneme_vocab_size = config["model"]["phoneme_vocab_size"]
|
||||
p_dropout = config["model"]["dropout"]
|
||||
EOS = config["model"]["EOS"]
|
||||
ffn_dim = hidden_dim * 4
|
||||
|
||||
self.n_layer = n_layer
|
||||
self.hidden_dim = hidden_dim
|
||||
self.n_head = n_head
|
||||
assert hidden_dim % n_head == 0
|
||||
self.head_dim = hidden_dim // n_head
|
||||
self.embedding_dim = embedding_dim
|
||||
self.vocab_size = vocab_size
|
||||
self.phoneme_vocab_size = phoneme_vocab_size
|
||||
self.p_dropout = p_dropout
|
||||
self.max_seq_length = max_seq_length
|
||||
self.max_batch_size = max_batch_size
|
||||
self.EOS = EOS
|
||||
assert self.EOS == self.vocab_size - 1
|
||||
|
||||
self.bert_proj = nn.Linear(1024, self.embedding_dim)
|
||||
self.ar_text_embedding = TokenEmbedding(
|
||||
self.embedding_dim, self.phoneme_vocab_size, self.p_dropout
|
||||
)
|
||||
self.ar_text_position = SinePositionalEmbedding(
|
||||
self.embedding_dim,
|
||||
dropout=0.1,
|
||||
scale=False,
|
||||
alpha=True,
|
||||
max_batch_size=max_batch_size,
|
||||
max_seq_len=max_seq_length,
|
||||
)
|
||||
self.ar_audio_embedding = TokenEmbedding(
|
||||
self.embedding_dim, self.vocab_size, self.p_dropout
|
||||
)
|
||||
self.ar_audio_position = SinePositionalEmbedding(
|
||||
self.embedding_dim,
|
||||
dropout=0.1,
|
||||
scale=False,
|
||||
alpha=True,
|
||||
max_batch_size=max_batch_size,
|
||||
max_seq_len=max_seq_length,
|
||||
)
|
||||
self.ar_predict_layer = nn.Linear(self.hidden_dim, self.vocab_size, bias=False)
|
||||
self.h = TransformerDecoder(
|
||||
hidden_dim,
|
||||
n_layer,
|
||||
n_head,
|
||||
ffn_dim,
|
||||
vocab_size,
|
||||
max_seq_length,
|
||||
max_batch_size,
|
||||
)
|
||||
|
||||
self._register_load_state_dict_pre_hook(self.load_hook)
|
||||
|
||||
def load_hook(self, state_dict, prefix, *args):
|
||||
model_keys = [key for key in state_dict if key.startswith("model.")]
|
||||
for key in model_keys:
|
||||
new_key = key[len("model.") :]
|
||||
state_dict[new_key] = state_dict.pop(key)
|
||||
|
||||
def init_cache(self, bsz: int = 0) -> nn.ModuleList:
|
||||
bsz = bsz or self.h.max_batch_size
|
||||
assert bsz <= self.h.max_batch_size
|
||||
seq_lens = self.h.max_seq_length
|
||||
device = self.bert_proj.bias.device
|
||||
dtype = self.bert_proj.bias.dtype
|
||||
return nn.ModuleList(
|
||||
[
|
||||
KVCacheNHD(bsz, seq_lens, self.n_head, self.head_dim)
|
||||
for _ in range(self.n_layer)
|
||||
],
|
||||
).to(device, dtype)
|
||||
|
||||
def embed(
|
||||
self,
|
||||
x: List[torch.Tensor],
|
||||
y: torch.Tensor,
|
||||
bert_features: List[torch.Tensor],
|
||||
):
|
||||
x_nested = torch.nested.nested_tensor(x)
|
||||
assert x_nested.size(0) <= self.max_batch_size
|
||||
bert_features_nested = torch.nested.nested_tensor(
|
||||
list(map(lambda t: t.transpose(0, 1), bert_features))
|
||||
)
|
||||
x_emb = self.ar_text_embedding.forward(x_nested)
|
||||
bert = self.bert_proj.forward(bert_features_nested)
|
||||
x_emb = x_emb + bert
|
||||
x_pos = self.ar_text_position.prefill(x_emb)
|
||||
|
||||
y_nested = torch.nested.nested_tensor(list(y.unbind(0)))
|
||||
y_emb = self.ar_audio_embedding.forward(y_nested)
|
||||
y_pos = self.ar_audio_position.prefill(y_emb)
|
||||
|
||||
xy_pos = torch.nested.nested_tensor(
|
||||
[torch.cat([x_pos[i], y_pos[i]]) for i in range(len(x))]
|
||||
)
|
||||
return xy_pos
|
||||
|
||||
def capture(
|
||||
self,
|
||||
input_pos: Tensor,
|
||||
x: Tensor,
|
||||
x_dec: Tensor,
|
||||
kv_caches,
|
||||
) -> CUDAGraph:
|
||||
s = torch.cuda.Stream()
|
||||
s.wait_stream(torch.cuda.current_stream())
|
||||
|
||||
graph = torch.cuda.CUDAGraph()
|
||||
|
||||
with torch.cuda.stream(s):
|
||||
for _ in range(5):
|
||||
self.h.forward(input_pos, x, kv_caches)
|
||||
torch.cuda.current_stream().wait_stream(s)
|
||||
|
||||
with torch.cuda.graph(graph):
|
||||
x_dec.copy_(self.h.forward(input_pos, x, kv_caches))
|
||||
torch.cuda.synchronize()
|
||||
|
||||
return graph
|
||||
|
||||
|
||||
# ─── CUDA Graph Runner ───────────────────────────────────────────────────────
|
||||
|
||||
|
||||
class CUDAGraphRunner:
|
||||
def __init__(
|
||||
self,
|
||||
decoder_model: T2SDecoder,
|
||||
device: torch.device = torch.device("cpu"),
|
||||
dtype: torch.dtype = torch.float32,
|
||||
) -> None:
|
||||
assert device.type in {"cpu", "cuda", "mps", "xpu", "mtia"}
|
||||
assert dtype in {torch.float16, torch.bfloat16, torch.float32}
|
||||
self.device = device
|
||||
self.dtype = dtype
|
||||
self.decoder_model: T2SDecoder = decoder_model.to(self.device, self.dtype)
|
||||
self.graph: Optional[CUDAGraph] = None
|
||||
self.xy_pos_ = torch.rand(
|
||||
(1, 1, decoder_model.embedding_dim), device=device
|
||||
).to(dtype)
|
||||
self.xy_dec_ = torch.rand(
|
||||
(1, 1, decoder_model.embedding_dim), device=device
|
||||
).to(dtype)
|
||||
self.kv_cache = decoder_model.init_cache(1)
|
||||
self.input_pos = torch.tensor([10]).int().cuda()
|
||||
|
||||
def _handle_request(self, request: T2SRequest):
|
||||
with self.device:
|
||||
for i in self.kv_cache:
|
||||
i.empty()
|
||||
|
||||
decoder = self.decoder_model
|
||||
session = T2SSession(decoder, request, device=self.device, dtype=self.dtype)
|
||||
self.input_pos.copy_(session.input_pos)
|
||||
|
||||
t1 = 0.0
|
||||
infer_speed = 0.0
|
||||
y = session.y
|
||||
bsz = y.size(0)
|
||||
|
||||
for idx in tqdm(range(1500)):
|
||||
if idx == 0:
|
||||
xy_dec = decoder.h.prefill(
|
||||
session.xy_pos, session.attn_mask_nested, self.kv_cache
|
||||
)
|
||||
xy_dec = torch.stack([t[[-1]] for t in xy_dec.unbind()])
|
||||
else:
|
||||
if (
|
||||
request.use_cuda_graph
|
||||
and self.graph is None
|
||||
and torch.cuda.is_available()
|
||||
):
|
||||
self.xy_pos_.copy_(session.xy_pos)
|
||||
self.graph = decoder.capture(
|
||||
self.input_pos,
|
||||
self.xy_pos_,
|
||||
self.xy_dec_,
|
||||
kv_caches=self.kv_cache,
|
||||
)
|
||||
|
||||
if self.graph:
|
||||
self.xy_pos_.copy_(session.xy_pos)
|
||||
self.graph.replay()
|
||||
xy_dec = self.xy_dec_.clone()
|
||||
else:
|
||||
xy_dec = decoder.h.forward(
|
||||
self.input_pos,
|
||||
session.xy_pos,
|
||||
self.kv_cache,
|
||||
)
|
||||
|
||||
logits = decoder.ar_predict_layer(xy_dec[:, -1])
|
||||
self.input_pos.add_(1)
|
||||
|
||||
if idx == 0:
|
||||
logits[:, -1] = float("-inf")
|
||||
|
||||
samples = session.sampler.sample(
|
||||
logits=logits,
|
||||
previous_tokens=session.y,
|
||||
top_k=request.top_k,
|
||||
top_p=request.top_p,
|
||||
repetition_penalty=request.repetition_penalty,
|
||||
temperature=request.temperature,
|
||||
)
|
||||
|
||||
session.y = torch.cat([session.y, samples], dim=1)
|
||||
|
||||
argmax_token = torch.argmax(logits, dim=-1)
|
||||
sample_token = samples.squeeze(1)
|
||||
EOS_mask = (argmax_token == decoder.EOS) | (
|
||||
sample_token == decoder.EOS
|
||||
)
|
||||
|
||||
newly_done_mask = EOS_mask & (~session.completed)
|
||||
newly_done_indices = newly_done_mask.nonzero()
|
||||
|
||||
if newly_done_indices.numel() > 0:
|
||||
session.y_results[newly_done_indices[0]] = session.y[
|
||||
newly_done_indices[0], session.y_len : -1
|
||||
].squeeze(0)
|
||||
session.completed[newly_done_indices] = True
|
||||
|
||||
if torch.all(session.completed).item():
|
||||
if session.y.size(1) == 0:
|
||||
session.y = torch.cat(
|
||||
[session.y, torch.zeros_like(samples)], dim=1
|
||||
)
|
||||
tqdm.write("Bad Zero Prediction")
|
||||
else:
|
||||
tqdm.write(
|
||||
f"T2S Decoding EOS {session.prefill_len.tolist().__str__().strip('[]')} -> \n"
|
||||
f"{[i.size(0) for i in session.y_results].__str__().strip('[]')}"
|
||||
)
|
||||
tqdm.write(
|
||||
f"Infer Speed: {(idx - 1) / (time.perf_counter() - t1):.2f} token/s"
|
||||
)
|
||||
infer_speed = (idx - 1) / (time.perf_counter() - t1)
|
||||
break
|
||||
|
||||
if (
|
||||
request.early_stop_num != -1
|
||||
and (session.y.size(1) - session.y_len) > request.early_stop_num
|
||||
) or idx == 1499:
|
||||
for i in range(bsz):
|
||||
if not session.completed[i].item():
|
||||
session.y_results[i] = session.y[i, session.y_len :]
|
||||
session.completed[i] = True
|
||||
break
|
||||
|
||||
y_emb = decoder.ar_audio_embedding(session.y[:, -1:])
|
||||
session.xy_pos = decoder.ar_audio_position.forward(
|
||||
self.input_pos - session.x_lens, y_emb
|
||||
)
|
||||
|
||||
if idx == 2:
|
||||
t1 = time.perf_counter()
|
||||
|
||||
if idx % 100 == 0 and self.device.type == "cuda":
|
||||
torch.cuda.empty_cache()
|
||||
|
||||
if self.device.type == "cuda":
|
||||
torch.cuda.empty_cache()
|
||||
|
||||
return session.y_results[: request.valid_length], infer_speed
|
||||
|
||||
def generate(self, request: T2SRequest) -> T2SResult:
|
||||
try:
|
||||
result, infer_speed = self._handle_request(request)
|
||||
t2s_result = T2SResult(
|
||||
result=result, infer_speed=infer_speed, status="Success"
|
||||
)
|
||||
except Exception as e:
|
||||
t2s_result = T2SResult(
|
||||
status="Error", exception=e, traceback=traceback.format_exc()
|
||||
)
|
||||
return t2s_result
|
||||
|
||||
@staticmethod
|
||||
def load_decoder(weights_path, max_batch_size=1) -> T2SDecoder:
|
||||
print(
|
||||
f"Loading Text2Semantic Weights from {weights_path} with CUDA Graph (SDPA) Implement"
|
||||
)
|
||||
dict_s1 = torch.load(
|
||||
weights_path, map_location="cpu", weights_only=False#, mmap=True
|
||||
)
|
||||
config = dict_s1["config"]
|
||||
decoder = T2SDecoder(config, max_batch_size=max_batch_size)
|
||||
state_dict = dict_s1["weight"]
|
||||
decoder.load_state_dict(state_dict)
|
||||
return decoder.eval()
|
||||
@@ -262,7 +262,7 @@ def make_reject_y(y_o, y_lens):
|
||||
reject_y = []
|
||||
reject_y_lens = []
|
||||
for b in range(bs):
|
||||
process_item_idx = torch.randint(0, 1, size=(1,))[0]
|
||||
process_item_idx = torch.randint(0, 2, size=(1,))[0]
|
||||
if process_item_idx == 0:
|
||||
new_y = repeat_P(y_o[b])
|
||||
reject_y.append(new_y)
|
||||
|
||||
@@ -354,7 +354,7 @@ class ScaledAdam(BatchedOptimizer):
|
||||
if ans < 1.0:
|
||||
first_state["num_clipped"] += 1
|
||||
if ans < 0.1:
|
||||
logging.warn(f"Scaling gradients by {ans}, model_norm_threshold={model_norm_threshold}")
|
||||
logging.warning(f"Scaling gradients by {ans}, model_norm_threshold={model_norm_threshold}")
|
||||
if self.show_dominant_parameters:
|
||||
assert p.shape[0] == len(param_names)
|
||||
self._show_gradient_dominating_parameter(tuples, tot_sumsq)
|
||||
@@ -362,7 +362,7 @@ class ScaledAdam(BatchedOptimizer):
|
||||
|
||||
def _show_gradient_dominating_parameter(self, tuples: List[Tuple[Tensor, dict, List[str]]], tot_sumsq: Tensor):
|
||||
"""
|
||||
Show information of parameter wihch dominanting tot_sumsq.
|
||||
Show information of parameter which dominating tot_sumsq.
|
||||
|
||||
Args:
|
||||
tuples: a list of tuples of (param, state, param_names)
|
||||
@@ -415,7 +415,7 @@ class ScaledAdam(BatchedOptimizer):
|
||||
dominant_grad,
|
||||
) = sorted_by_proportion[dominant_param_name]
|
||||
logging.info(
|
||||
f"Parameter Dominanting tot_sumsq {dominant_param_name}"
|
||||
f"Parameter Dominating tot_sumsq {dominant_param_name}"
|
||||
f" with proportion {dominant_proportion:.2f},"
|
||||
f" where dominant_sumsq=(grad_sumsq*orig_rms_sq)"
|
||||
f"={dominant_sumsq:.3e},"
|
||||
|
||||
@@ -8,30 +8,30 @@ def multi_head_attention_forward_patched(
|
||||
query,
|
||||
key,
|
||||
value,
|
||||
embed_dim_to_check: int,
|
||||
num_heads: int,
|
||||
embed_dim_to_check,
|
||||
num_heads,
|
||||
in_proj_weight,
|
||||
in_proj_bias: Optional[Tensor],
|
||||
bias_k: Optional[Tensor],
|
||||
bias_v: Optional[Tensor],
|
||||
add_zero_attn: bool,
|
||||
dropout_p: float,
|
||||
out_proj_weight: Tensor,
|
||||
out_proj_bias: Optional[Tensor],
|
||||
training: bool = True,
|
||||
key_padding_mask: Optional[Tensor] = None,
|
||||
need_weights: bool = True,
|
||||
attn_mask: Optional[Tensor] = None,
|
||||
use_separate_proj_weight: bool = False,
|
||||
q_proj_weight: Optional[Tensor] = None,
|
||||
k_proj_weight: Optional[Tensor] = None,
|
||||
v_proj_weight: Optional[Tensor] = None,
|
||||
static_k: Optional[Tensor] = None,
|
||||
static_v: Optional[Tensor] = None,
|
||||
average_attn_weights: bool = True,
|
||||
is_causal: bool = False,
|
||||
in_proj_bias,
|
||||
bias_k,
|
||||
bias_v,
|
||||
add_zero_attn,
|
||||
dropout_p,
|
||||
out_proj_weight,
|
||||
out_proj_bias,
|
||||
training=True,
|
||||
key_padding_mask=None,
|
||||
need_weights=True,
|
||||
attn_mask=None,
|
||||
use_separate_proj_weight=False,
|
||||
q_proj_weight=None,
|
||||
k_proj_weight=None,
|
||||
v_proj_weight=None,
|
||||
static_k=None,
|
||||
static_v=None,
|
||||
average_attn_weights=True,
|
||||
is_causal=False,
|
||||
cache=None,
|
||||
) -> Tuple[Tensor, Optional[Tensor]]:
|
||||
):
|
||||
# set up shape vars
|
||||
_, _, embed_dim = query.shape
|
||||
attn_mask = _canonical_mask(
|
||||
|
||||
@@ -1,69 +0,0 @@
|
||||
# Copyright (c) 2024 NVIDIA CORPORATION.
|
||||
# Licensed under the MIT license.
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
from alias_free_activation.torch.resample import UpSample1d, DownSample1d
|
||||
|
||||
# load fused CUDA kernel: this enables importing anti_alias_activation_cuda
|
||||
from alias_free_activation.cuda import load
|
||||
|
||||
anti_alias_activation_cuda = load.load()
|
||||
|
||||
|
||||
class FusedAntiAliasActivation(torch.autograd.Function):
|
||||
"""
|
||||
Assumes filter size 12, replication padding on upsampling/downsampling, and logscale alpha/beta parameters as inputs.
|
||||
The hyperparameters are hard-coded in the kernel to maximize speed.
|
||||
NOTE: The fused kenrel is incorrect for Activation1d with different hyperparameters.
|
||||
"""
|
||||
|
||||
@staticmethod
|
||||
def forward(ctx, inputs, up_ftr, down_ftr, alpha, beta):
|
||||
activation_results = anti_alias_activation_cuda.forward(inputs, up_ftr, down_ftr, alpha, beta)
|
||||
|
||||
return activation_results
|
||||
|
||||
@staticmethod
|
||||
def backward(ctx, output_grads):
|
||||
raise NotImplementedError
|
||||
return output_grads, None, None
|
||||
|
||||
|
||||
class Activation1d(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
activation,
|
||||
up_ratio: int = 2,
|
||||
down_ratio: int = 2,
|
||||
up_kernel_size: int = 12,
|
||||
down_kernel_size: int = 12,
|
||||
fused: bool = True,
|
||||
):
|
||||
super().__init__()
|
||||
self.up_ratio = up_ratio
|
||||
self.down_ratio = down_ratio
|
||||
self.act = activation
|
||||
self.upsample = UpSample1d(up_ratio, up_kernel_size)
|
||||
self.downsample = DownSample1d(down_ratio, down_kernel_size)
|
||||
|
||||
self.fused = fused # Whether to use fused CUDA kernel or not
|
||||
|
||||
def forward(self, x):
|
||||
if not self.fused:
|
||||
x = self.upsample(x)
|
||||
x = self.act(x)
|
||||
x = self.downsample(x)
|
||||
return x
|
||||
else:
|
||||
if self.act.__class__.__name__ == "Snake":
|
||||
beta = self.act.alpha.data # Snake uses same params for alpha and beta
|
||||
else:
|
||||
beta = self.act.beta.data # Snakebeta uses different params for alpha and beta
|
||||
alpha = self.act.alpha.data
|
||||
if not self.act.alpha_logscale: # Exp baked into cuda kernel, cancel it out with a log
|
||||
alpha = torch.log(alpha)
|
||||
beta = torch.log(beta)
|
||||
|
||||
x = FusedAntiAliasActivation.apply(x, self.upsample.filter, self.downsample.lowpass.filter, alpha, beta)
|
||||
return x
|
||||
@@ -1,23 +0,0 @@
|
||||
/* coding=utf-8
|
||||
* Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved.
|
||||
*
|
||||
* Licensed under the Apache License, Version 2.0 (the "License");
|
||||
* you may not use this file except in compliance with the License.
|
||||
* You may obtain a copy of the License at
|
||||
*
|
||||
* http://www.apache.org/licenses/LICENSE-2.0
|
||||
*
|
||||
* Unless required by applicable law or agreed to in writing, software
|
||||
* distributed under the License is distributed on an "AS IS" BASIS,
|
||||
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
* See the License for the specific language governing permissions and
|
||||
* limitations under the License.
|
||||
*/
|
||||
|
||||
#include <torch/extension.h>
|
||||
|
||||
extern "C" torch::Tensor fwd_cuda(torch::Tensor const &input, torch::Tensor const &up_filter, torch::Tensor const &down_filter, torch::Tensor const &alpha, torch::Tensor const &beta);
|
||||
|
||||
PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
|
||||
m.def("forward", &fwd_cuda, "Anti-Alias Activation forward (CUDA)");
|
||||
}
|
||||
@@ -1,246 +0,0 @@
|
||||
/* coding=utf-8
|
||||
* Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved.
|
||||
*
|
||||
* Licensed under the Apache License, Version 2.0 (the "License");
|
||||
* you may not use this file except in compliance with the License.
|
||||
* You may obtain a copy of the License at
|
||||
*
|
||||
* http://www.apache.org/licenses/LICENSE-2.0
|
||||
*
|
||||
* Unless required by applicable law or agreed to in writing, software
|
||||
* distributed under the License is distributed on an "AS IS" BASIS,
|
||||
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
* See the License for the specific language governing permissions and
|
||||
* limitations under the License.
|
||||
*/
|
||||
|
||||
#include <ATen/ATen.h>
|
||||
#include <cuda.h>
|
||||
#include <cuda_runtime.h>
|
||||
#include <cuda_fp16.h>
|
||||
#include <cuda_profiler_api.h>
|
||||
#include <ATen/cuda/CUDAContext.h>
|
||||
#include <torch/extension.h>
|
||||
#include "type_shim.h"
|
||||
#include <assert.h>
|
||||
#include <cfloat>
|
||||
#include <limits>
|
||||
#include <stdint.h>
|
||||
#include <c10/macros/Macros.h>
|
||||
|
||||
namespace
|
||||
{
|
||||
// Hard-coded hyperparameters
|
||||
// WARP_SIZE and WARP_BATCH must match the return values batches_per_warp and
|
||||
constexpr int ELEMENTS_PER_LDG_STG = 1; //(WARP_ITERATIONS < 4) ? 1 : 4;
|
||||
constexpr int BUFFER_SIZE = 32;
|
||||
constexpr int FILTER_SIZE = 12;
|
||||
constexpr int HALF_FILTER_SIZE = 6;
|
||||
constexpr int UPSAMPLE_REPLICATION_PAD = 5; // 5 on each side, matching torch impl
|
||||
constexpr int DOWNSAMPLE_REPLICATION_PAD_LEFT = 5; // matching torch impl
|
||||
constexpr int DOWNSAMPLE_REPLICATION_PAD_RIGHT = 6; // matching torch impl
|
||||
|
||||
template <typename input_t, typename output_t, typename acc_t>
|
||||
__global__ void anti_alias_activation_forward(
|
||||
output_t *dst,
|
||||
const input_t *src,
|
||||
const input_t *up_ftr,
|
||||
const input_t *down_ftr,
|
||||
const input_t *alpha,
|
||||
const input_t *beta,
|
||||
int batch_size,
|
||||
int channels,
|
||||
int seq_len)
|
||||
{
|
||||
// Up and downsample filters
|
||||
input_t up_filter[FILTER_SIZE];
|
||||
input_t down_filter[FILTER_SIZE];
|
||||
|
||||
// Load data from global memory including extra indices reserved for replication paddings
|
||||
input_t elements[2 * FILTER_SIZE + 2 * BUFFER_SIZE + 2 * UPSAMPLE_REPLICATION_PAD] = {0};
|
||||
input_t intermediates[2 * FILTER_SIZE + 2 * BUFFER_SIZE + DOWNSAMPLE_REPLICATION_PAD_LEFT + DOWNSAMPLE_REPLICATION_PAD_RIGHT] = {0};
|
||||
|
||||
// Output stores downsampled output before writing to dst
|
||||
output_t output[BUFFER_SIZE];
|
||||
|
||||
// blockDim/threadIdx = (128, 1, 1)
|
||||
// gridDim/blockIdx = (seq_blocks, channels, batches)
|
||||
int block_offset = (blockIdx.x * 128 * BUFFER_SIZE + seq_len * (blockIdx.y + gridDim.y * blockIdx.z));
|
||||
int local_offset = threadIdx.x * BUFFER_SIZE;
|
||||
int seq_offset = blockIdx.x * 128 * BUFFER_SIZE + local_offset;
|
||||
|
||||
// intermediate have double the seq_len
|
||||
int intermediate_local_offset = threadIdx.x * BUFFER_SIZE * 2;
|
||||
int intermediate_seq_offset = blockIdx.x * 128 * BUFFER_SIZE * 2 + intermediate_local_offset;
|
||||
|
||||
// Get values needed for replication padding before moving pointer
|
||||
const input_t *right_most_pntr = src + (seq_len * (blockIdx.y + gridDim.y * blockIdx.z));
|
||||
input_t seq_left_most_value = right_most_pntr[0];
|
||||
input_t seq_right_most_value = right_most_pntr[seq_len - 1];
|
||||
|
||||
// Move src and dst pointers
|
||||
src += block_offset + local_offset;
|
||||
dst += block_offset + local_offset;
|
||||
|
||||
// Alpha and beta values for snake activatons. Applies exp by default
|
||||
alpha = alpha + blockIdx.y;
|
||||
input_t alpha_val = expf(alpha[0]);
|
||||
beta = beta + blockIdx.y;
|
||||
input_t beta_val = expf(beta[0]);
|
||||
|
||||
#pragma unroll
|
||||
for (int it = 0; it < FILTER_SIZE; it += 1)
|
||||
{
|
||||
up_filter[it] = up_ftr[it];
|
||||
down_filter[it] = down_ftr[it];
|
||||
}
|
||||
|
||||
// Apply replication padding for upsampling, matching torch impl
|
||||
#pragma unroll
|
||||
for (int it = -HALF_FILTER_SIZE; it < BUFFER_SIZE + HALF_FILTER_SIZE; it += 1)
|
||||
{
|
||||
int element_index = seq_offset + it; // index for element
|
||||
if ((element_index < 0) && (element_index >= -UPSAMPLE_REPLICATION_PAD))
|
||||
{
|
||||
elements[2 * (HALF_FILTER_SIZE + it)] = 2 * seq_left_most_value;
|
||||
}
|
||||
if ((element_index >= seq_len) && (element_index < seq_len + UPSAMPLE_REPLICATION_PAD))
|
||||
{
|
||||
elements[2 * (HALF_FILTER_SIZE + it)] = 2 * seq_right_most_value;
|
||||
}
|
||||
if ((element_index >= 0) && (element_index < seq_len))
|
||||
{
|
||||
elements[2 * (HALF_FILTER_SIZE + it)] = 2 * src[it];
|
||||
}
|
||||
}
|
||||
|
||||
// Apply upsampling strided convolution and write to intermediates. It reserves DOWNSAMPLE_REPLICATION_PAD_LEFT for replication padding of the downsampilng conv later
|
||||
#pragma unroll
|
||||
for (int it = 0; it < (2 * BUFFER_SIZE + 2 * FILTER_SIZE); it += 1)
|
||||
{
|
||||
input_t acc = 0.0;
|
||||
int element_index = intermediate_seq_offset + it; // index for intermediate
|
||||
#pragma unroll
|
||||
for (int f_idx = 0; f_idx < FILTER_SIZE; f_idx += 1)
|
||||
{
|
||||
if ((element_index + f_idx) >= 0)
|
||||
{
|
||||
acc += up_filter[f_idx] * elements[it + f_idx];
|
||||
}
|
||||
}
|
||||
intermediates[it + DOWNSAMPLE_REPLICATION_PAD_LEFT] = acc;
|
||||
}
|
||||
|
||||
// Apply activation function. It reserves DOWNSAMPLE_REPLICATION_PAD_LEFT and DOWNSAMPLE_REPLICATION_PAD_RIGHT for replication padding of the downsampilng conv later
|
||||
double no_div_by_zero = 0.000000001;
|
||||
#pragma unroll
|
||||
for (int it = 0; it < 2 * BUFFER_SIZE + 2 * FILTER_SIZE; it += 1)
|
||||
{
|
||||
intermediates[it + DOWNSAMPLE_REPLICATION_PAD_LEFT] += (1.0 / (beta_val + no_div_by_zero)) * sinf(intermediates[it + DOWNSAMPLE_REPLICATION_PAD_LEFT] * alpha_val) * sinf(intermediates[it + DOWNSAMPLE_REPLICATION_PAD_LEFT] * alpha_val);
|
||||
}
|
||||
|
||||
// Apply replication padding before downsampling conv from intermediates
|
||||
#pragma unroll
|
||||
for (int it = 0; it < DOWNSAMPLE_REPLICATION_PAD_LEFT; it += 1)
|
||||
{
|
||||
intermediates[it] = intermediates[DOWNSAMPLE_REPLICATION_PAD_LEFT];
|
||||
}
|
||||
#pragma unroll
|
||||
for (int it = DOWNSAMPLE_REPLICATION_PAD_LEFT + 2 * BUFFER_SIZE + 2 * FILTER_SIZE; it < DOWNSAMPLE_REPLICATION_PAD_LEFT + 2 * BUFFER_SIZE + 2 * FILTER_SIZE + DOWNSAMPLE_REPLICATION_PAD_RIGHT; it += 1)
|
||||
{
|
||||
intermediates[it] = intermediates[DOWNSAMPLE_REPLICATION_PAD_LEFT + 2 * BUFFER_SIZE + 2 * FILTER_SIZE - 1];
|
||||
}
|
||||
|
||||
// Apply downsample strided convolution (assuming stride=2) from intermediates
|
||||
#pragma unroll
|
||||
for (int it = 0; it < BUFFER_SIZE; it += 1)
|
||||
{
|
||||
input_t acc = 0.0;
|
||||
#pragma unroll
|
||||
for (int f_idx = 0; f_idx < FILTER_SIZE; f_idx += 1)
|
||||
{
|
||||
// Add constant DOWNSAMPLE_REPLICATION_PAD_RIGHT to match torch implementation
|
||||
acc += down_filter[f_idx] * intermediates[it * 2 + f_idx + DOWNSAMPLE_REPLICATION_PAD_RIGHT];
|
||||
}
|
||||
output[it] = acc;
|
||||
}
|
||||
|
||||
// Write output to dst
|
||||
#pragma unroll
|
||||
for (int it = 0; it < BUFFER_SIZE; it += ELEMENTS_PER_LDG_STG)
|
||||
{
|
||||
int element_index = seq_offset + it;
|
||||
if (element_index < seq_len)
|
||||
{
|
||||
dst[it] = output[it];
|
||||
}
|
||||
}
|
||||
|
||||
}
|
||||
|
||||
template <typename input_t, typename output_t, typename acc_t>
|
||||
void dispatch_anti_alias_activation_forward(
|
||||
output_t *dst,
|
||||
const input_t *src,
|
||||
const input_t *up_ftr,
|
||||
const input_t *down_ftr,
|
||||
const input_t *alpha,
|
||||
const input_t *beta,
|
||||
int batch_size,
|
||||
int channels,
|
||||
int seq_len)
|
||||
{
|
||||
if (seq_len == 0)
|
||||
{
|
||||
return;
|
||||
}
|
||||
else
|
||||
{
|
||||
// Use 128 threads per block to maximimize gpu utilization
|
||||
constexpr int threads_per_block = 128;
|
||||
constexpr int seq_len_per_block = 4096;
|
||||
int blocks_per_seq_len = (seq_len + seq_len_per_block - 1) / seq_len_per_block;
|
||||
dim3 blocks(blocks_per_seq_len, channels, batch_size);
|
||||
dim3 threads(threads_per_block, 1, 1);
|
||||
|
||||
anti_alias_activation_forward<input_t, output_t, acc_t>
|
||||
<<<blocks, threads, 0, at::cuda::getCurrentCUDAStream()>>>(dst, src, up_ftr, down_ftr, alpha, beta, batch_size, channels, seq_len);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
extern "C" torch::Tensor fwd_cuda(torch::Tensor const &input, torch::Tensor const &up_filter, torch::Tensor const &down_filter, torch::Tensor const &alpha, torch::Tensor const &beta)
|
||||
{
|
||||
// Input is a 3d tensor with dimensions [batches, channels, seq_len]
|
||||
const int batches = input.size(0);
|
||||
const int channels = input.size(1);
|
||||
const int seq_len = input.size(2);
|
||||
|
||||
// Output
|
||||
auto act_options = input.options().requires_grad(false);
|
||||
|
||||
torch::Tensor anti_alias_activation_results =
|
||||
torch::empty({batches, channels, seq_len}, act_options);
|
||||
|
||||
void *input_ptr = static_cast<void *>(input.data_ptr());
|
||||
void *up_filter_ptr = static_cast<void *>(up_filter.data_ptr());
|
||||
void *down_filter_ptr = static_cast<void *>(down_filter.data_ptr());
|
||||
void *alpha_ptr = static_cast<void *>(alpha.data_ptr());
|
||||
void *beta_ptr = static_cast<void *>(beta.data_ptr());
|
||||
void *anti_alias_activation_results_ptr = static_cast<void *>(anti_alias_activation_results.data_ptr());
|
||||
|
||||
DISPATCH_FLOAT_HALF_AND_BFLOAT(
|
||||
input.scalar_type(),
|
||||
"dispatch anti alias activation_forward",
|
||||
dispatch_anti_alias_activation_forward<scalar_t, scalar_t, float>(
|
||||
reinterpret_cast<scalar_t *>(anti_alias_activation_results_ptr),
|
||||
reinterpret_cast<const scalar_t *>(input_ptr),
|
||||
reinterpret_cast<const scalar_t *>(up_filter_ptr),
|
||||
reinterpret_cast<const scalar_t *>(down_filter_ptr),
|
||||
reinterpret_cast<const scalar_t *>(alpha_ptr),
|
||||
reinterpret_cast<const scalar_t *>(beta_ptr),
|
||||
batches,
|
||||
channels,
|
||||
seq_len););
|
||||
return anti_alias_activation_results;
|
||||
}
|
||||
@@ -1 +0,0 @@
|
||||
|
||||
@@ -1,29 +0,0 @@
|
||||
/* coding=utf-8
|
||||
* Copyright (c) 2020, NVIDIA CORPORATION. All rights reserved.
|
||||
*
|
||||
* Licensed under the Apache License, Version 2.0 (the "License");
|
||||
* you may not use this file except in compliance with the License.
|
||||
* You may obtain a copy of the License at
|
||||
*
|
||||
* http://www.apache.org/licenses/LICENSE-2.0
|
||||
*
|
||||
* Unless required by applicable law or agreed to in writing, software
|
||||
* distributed under the License is distributed on an "AS IS" BASIS,
|
||||
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
* See the License for the specific language governing permissions and
|
||||
* limitations under the License.
|
||||
*/
|
||||
|
||||
/*This code is copied fron NVIDIA apex:
|
||||
* https://github.com/NVIDIA/apex
|
||||
* with minor changes. */
|
||||
|
||||
#ifndef TORCH_CHECK
|
||||
#define TORCH_CHECK AT_CHECK
|
||||
#endif
|
||||
|
||||
#ifdef VERSION_GE_1_3
|
||||
#define DATA_PTR data_ptr
|
||||
#else
|
||||
#define DATA_PTR data
|
||||
#endif
|
||||
@@ -1,82 +0,0 @@
|
||||
# Copyright (c) 2024 NVIDIA CORPORATION.
|
||||
# Licensed under the MIT license.
|
||||
|
||||
import os
|
||||
import pathlib
|
||||
import subprocess
|
||||
|
||||
from torch.utils import cpp_extension
|
||||
|
||||
"""
|
||||
Setting this param to a list has a problem of generating different compilation commands (with diferent order of architectures) and leading to recompilation of fused kernels.
|
||||
Set it to empty stringo avoid recompilation and assign arch flags explicity in extra_cuda_cflags below
|
||||
"""
|
||||
os.environ["TORCH_CUDA_ARCH_LIST"] = ""
|
||||
|
||||
|
||||
def load():
|
||||
# Check if cuda 11 is installed for compute capability 8.0
|
||||
cc_flag = []
|
||||
_, bare_metal_major, _ = _get_cuda_bare_metal_version(cpp_extension.CUDA_HOME)
|
||||
if int(bare_metal_major) >= 11:
|
||||
cc_flag.append("-gencode")
|
||||
cc_flag.append("arch=compute_80,code=sm_80")
|
||||
|
||||
# Build path
|
||||
srcpath = pathlib.Path(__file__).parent.absolute()
|
||||
buildpath = srcpath / "build"
|
||||
_create_build_dir(buildpath)
|
||||
|
||||
# Helper function to build the kernels.
|
||||
def _cpp_extention_load_helper(name, sources, extra_cuda_flags):
|
||||
return cpp_extension.load(
|
||||
name=name,
|
||||
sources=sources,
|
||||
build_directory=buildpath,
|
||||
extra_cflags=[
|
||||
"-O3",
|
||||
],
|
||||
extra_cuda_cflags=[
|
||||
"-O3",
|
||||
"-gencode",
|
||||
"arch=compute_70,code=sm_70",
|
||||
"--use_fast_math",
|
||||
]
|
||||
+ extra_cuda_flags
|
||||
+ cc_flag,
|
||||
verbose=True,
|
||||
)
|
||||
|
||||
extra_cuda_flags = [
|
||||
"-U__CUDA_NO_HALF_OPERATORS__",
|
||||
"-U__CUDA_NO_HALF_CONVERSIONS__",
|
||||
"--expt-relaxed-constexpr",
|
||||
"--expt-extended-lambda",
|
||||
]
|
||||
|
||||
sources = [
|
||||
srcpath / "anti_alias_activation.cpp",
|
||||
srcpath / "anti_alias_activation_cuda.cu",
|
||||
]
|
||||
anti_alias_activation_cuda = _cpp_extention_load_helper("anti_alias_activation_cuda", sources, extra_cuda_flags)
|
||||
|
||||
return anti_alias_activation_cuda
|
||||
|
||||
|
||||
def _get_cuda_bare_metal_version(cuda_dir):
|
||||
raw_output = subprocess.check_output([cuda_dir + "/bin/nvcc", "-V"], universal_newlines=True)
|
||||
output = raw_output.split()
|
||||
release_idx = output.index("release") + 1
|
||||
release = output[release_idx].split(".")
|
||||
bare_metal_major = release[0]
|
||||
bare_metal_minor = release[1][0]
|
||||
|
||||
return raw_output, bare_metal_major, bare_metal_minor
|
||||
|
||||
|
||||
def _create_build_dir(buildpath):
|
||||
try:
|
||||
os.mkdir(buildpath)
|
||||
except OSError:
|
||||
if not os.path.isdir(buildpath):
|
||||
print(f"Creation of the build directory {buildpath} failed")
|
||||
@@ -1,92 +0,0 @@
|
||||
/* coding=utf-8
|
||||
* Copyright (c) 2020, NVIDIA CORPORATION. All rights reserved.
|
||||
*
|
||||
* Licensed under the Apache License, Version 2.0 (the "License");
|
||||
* you may not use this file except in compliance with the License.
|
||||
* You may obtain a copy of the License at
|
||||
*
|
||||
* http://www.apache.org/licenses/LICENSE-2.0
|
||||
*
|
||||
* Unless required by applicable law or agreed to in writing, software
|
||||
* distributed under the License is distributed on an "AS IS" BASIS,
|
||||
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
* See the License for the specific language governing permissions and
|
||||
* limitations under the License.
|
||||
*/
|
||||
|
||||
#include <ATen/ATen.h>
|
||||
#include "compat.h"
|
||||
|
||||
#define DISPATCH_FLOAT_HALF_AND_BFLOAT(TYPE, NAME, ...) \
|
||||
switch (TYPE) \
|
||||
{ \
|
||||
case at::ScalarType::Float: \
|
||||
{ \
|
||||
using scalar_t = float; \
|
||||
__VA_ARGS__; \
|
||||
break; \
|
||||
} \
|
||||
case at::ScalarType::Half: \
|
||||
{ \
|
||||
using scalar_t = at::Half; \
|
||||
__VA_ARGS__; \
|
||||
break; \
|
||||
} \
|
||||
case at::ScalarType::BFloat16: \
|
||||
{ \
|
||||
using scalar_t = at::BFloat16; \
|
||||
__VA_ARGS__; \
|
||||
break; \
|
||||
} \
|
||||
default: \
|
||||
AT_ERROR(#NAME, " not implemented for '", toString(TYPE), "'"); \
|
||||
}
|
||||
|
||||
#define DISPATCH_FLOAT_HALF_AND_BFLOAT_INOUT_TYPES(TYPEIN, TYPEOUT, NAME, ...) \
|
||||
switch (TYPEIN) \
|
||||
{ \
|
||||
case at::ScalarType::Float: \
|
||||
{ \
|
||||
using scalar_t_in = float; \
|
||||
switch (TYPEOUT) \
|
||||
{ \
|
||||
case at::ScalarType::Float: \
|
||||
{ \
|
||||
using scalar_t_out = float; \
|
||||
__VA_ARGS__; \
|
||||
break; \
|
||||
} \
|
||||
case at::ScalarType::Half: \
|
||||
{ \
|
||||
using scalar_t_out = at::Half; \
|
||||
__VA_ARGS__; \
|
||||
break; \
|
||||
} \
|
||||
case at::ScalarType::BFloat16: \
|
||||
{ \
|
||||
using scalar_t_out = at::BFloat16; \
|
||||
__VA_ARGS__; \
|
||||
break; \
|
||||
} \
|
||||
default: \
|
||||
AT_ERROR(#NAME, " not implemented for '", toString(TYPEOUT), "'"); \
|
||||
} \
|
||||
break; \
|
||||
} \
|
||||
case at::ScalarType::Half: \
|
||||
{ \
|
||||
using scalar_t_in = at::Half; \
|
||||
using scalar_t_out = at::Half; \
|
||||
__VA_ARGS__; \
|
||||
break; \
|
||||
} \
|
||||
case at::ScalarType::BFloat16: \
|
||||
{ \
|
||||
using scalar_t_in = at::BFloat16; \
|
||||
using scalar_t_out = at::BFloat16; \
|
||||
__VA_ARGS__; \
|
||||
break; \
|
||||
} \
|
||||
default: \
|
||||
AT_ERROR(#NAME, " not implemented for '", toString(TYPEIN), "'"); \
|
||||
}
|
||||
@@ -89,28 +89,18 @@ class AMPBlock1(torch.nn.Module):
|
||||
|
||||
self.num_layers = len(self.convs1) + len(self.convs2) # Total number of conv layers
|
||||
|
||||
# Select which Activation1d, lazy-load cuda version to ensure backward compatibility
|
||||
if self.h.get("use_cuda_kernel", False):
|
||||
from .alias_free_activation.cuda.activation1d import (
|
||||
Activation1d as CudaActivation1d,
|
||||
)
|
||||
|
||||
Activation1d = CudaActivation1d
|
||||
else:
|
||||
Activation1d = TorchActivation1d
|
||||
|
||||
# Activation functions
|
||||
if activation == "snake":
|
||||
self.activations = nn.ModuleList(
|
||||
[
|
||||
Activation1d(activation=activations.Snake(channels, alpha_logscale=h.snake_logscale))
|
||||
TorchActivation1d(activation=activations.Snake(channels, alpha_logscale=h.snake_logscale))
|
||||
for _ in range(self.num_layers)
|
||||
]
|
||||
)
|
||||
elif activation == "snakebeta":
|
||||
self.activations = nn.ModuleList(
|
||||
[
|
||||
Activation1d(activation=activations.SnakeBeta(channels, alpha_logscale=h.snake_logscale))
|
||||
TorchActivation1d(activation=activations.SnakeBeta(channels, alpha_logscale=h.snake_logscale))
|
||||
for _ in range(self.num_layers)
|
||||
]
|
||||
)
|
||||
@@ -181,28 +171,18 @@ class AMPBlock2(torch.nn.Module):
|
||||
|
||||
self.num_layers = len(self.convs) # Total number of conv layers
|
||||
|
||||
# Select which Activation1d, lazy-load cuda version to ensure backward compatibility
|
||||
if self.h.get("use_cuda_kernel", False):
|
||||
from .alias_free_activation.cuda.activation1d import (
|
||||
Activation1d as CudaActivation1d,
|
||||
)
|
||||
|
||||
Activation1d = CudaActivation1d
|
||||
else:
|
||||
Activation1d = TorchActivation1d
|
||||
|
||||
# Activation functions
|
||||
if activation == "snake":
|
||||
self.activations = nn.ModuleList(
|
||||
[
|
||||
Activation1d(activation=activations.Snake(channels, alpha_logscale=h.snake_logscale))
|
||||
TorchActivation1d(activation=activations.Snake(channels, alpha_logscale=h.snake_logscale))
|
||||
for _ in range(self.num_layers)
|
||||
]
|
||||
)
|
||||
elif activation == "snakebeta":
|
||||
self.activations = nn.ModuleList(
|
||||
[
|
||||
Activation1d(activation=activations.SnakeBeta(channels, alpha_logscale=h.snake_logscale))
|
||||
TorchActivation1d(activation=activations.SnakeBeta(channels, alpha_logscale=h.snake_logscale))
|
||||
for _ in range(self.num_layers)
|
||||
]
|
||||
)
|
||||
@@ -235,31 +215,17 @@ class BigVGAN(
|
||||
):
|
||||
"""
|
||||
BigVGAN is a neural vocoder model that applies anti-aliased periodic activation for residual blocks (resblocks).
|
||||
New in BigVGAN-v2: it can optionally use optimized CUDA kernels for AMP (anti-aliased multi-periodicity) blocks.
|
||||
|
||||
Args:
|
||||
h (AttrDict): Hyperparameters.
|
||||
use_cuda_kernel (bool): If set to True, loads optimized CUDA kernels for AMP. This should be used for inference only, as training is not supported with CUDA kernels.
|
||||
|
||||
Note:
|
||||
- The `use_cuda_kernel` parameter should be used for inference only, as training with CUDA kernels is not supported.
|
||||
- Ensure that the activation function is correctly specified in the hyperparameters (h.activation).
|
||||
"""
|
||||
|
||||
def __init__(self, h: AttrDict, use_cuda_kernel: bool = False):
|
||||
def __init__(self, h: AttrDict):
|
||||
super().__init__()
|
||||
self.h = h
|
||||
self.h["use_cuda_kernel"] = use_cuda_kernel
|
||||
|
||||
# Select which Activation1d, lazy-load cuda version to ensure backward compatibility
|
||||
if self.h.get("use_cuda_kernel", False):
|
||||
from .alias_free_activation.cuda.activation1d import (
|
||||
Activation1d as CudaActivation1d,
|
||||
)
|
||||
|
||||
Activation1d = CudaActivation1d
|
||||
else:
|
||||
Activation1d = TorchActivation1d
|
||||
|
||||
self.num_kernels = len(h.resblock_kernel_sizes)
|
||||
self.num_upsamples = len(h.upsample_rates)
|
||||
@@ -312,7 +278,7 @@ class BigVGAN(
|
||||
"activation incorrectly specified. check the config file and look for 'activation'."
|
||||
)
|
||||
|
||||
self.activation_post = Activation1d(activation=activation_post)
|
||||
self.activation_post = TorchActivation1d(activation=activation_post)
|
||||
|
||||
# Whether to use bias for the final conv_post. Default to True for backward compatibility
|
||||
self.use_bias_at_final = h.get("use_bias_at_final", True)
|
||||
@@ -393,7 +359,6 @@ class BigVGAN(
|
||||
token: Union[str, bool, None],
|
||||
map_location: str = "cpu", # Additional argument
|
||||
strict: bool = False, # Additional argument
|
||||
use_cuda_kernel: bool = False,
|
||||
**model_kwargs,
|
||||
):
|
||||
"""Load Pytorch pretrained weights and return the loaded model."""
|
||||
@@ -417,17 +382,7 @@ class BigVGAN(
|
||||
h = load_hparams_from_json(config_file)
|
||||
|
||||
# instantiate BigVGAN using h
|
||||
if use_cuda_kernel:
|
||||
print(
|
||||
"[WARNING] You have specified use_cuda_kernel=True during BigVGAN.from_pretrained(). Only inference is supported (training is not implemented)!"
|
||||
)
|
||||
print(
|
||||
"[WARNING] You need nvcc and ninja installed in your system that matches your PyTorch build is using to build the kernel. If not, the model will fail to initialize or generate incorrect waveform!"
|
||||
)
|
||||
print(
|
||||
"[WARNING] For detail, see the official GitHub repository: https://github.com/NVIDIA/BigVGAN?tab=readme-ov-file#using-custom-cuda-kernel-for-synthesis"
|
||||
)
|
||||
model = cls(h, use_cuda_kernel=use_cuda_kernel)
|
||||
model = cls(h)
|
||||
|
||||
# Download and load pretrained generator weight
|
||||
if os.path.isdir(model_id):
|
||||
|
||||
@@ -1,45 +0,0 @@
|
||||
{
|
||||
"resblock": "1",
|
||||
"num_gpus": 0,
|
||||
"batch_size": 32,
|
||||
"learning_rate": 0.0001,
|
||||
"adam_b1": 0.8,
|
||||
"adam_b2": 0.99,
|
||||
"lr_decay": 0.9999996,
|
||||
"seed": 1234,
|
||||
|
||||
"upsample_rates": [4,4,2,2,2,2],
|
||||
"upsample_kernel_sizes": [8,8,4,4,4,4],
|
||||
"upsample_initial_channel": 1536,
|
||||
"resblock_kernel_sizes": [3,7,11],
|
||||
"resblock_dilation_sizes": [[1,3,5], [1,3,5], [1,3,5]],
|
||||
|
||||
"activation": "snakebeta",
|
||||
"snake_logscale": true,
|
||||
|
||||
"resolutions": [[1024, 120, 600], [2048, 240, 1200], [512, 50, 240]],
|
||||
"mpd_reshapes": [2, 3, 5, 7, 11],
|
||||
"use_spectral_norm": false,
|
||||
"discriminator_channel_mult": 1,
|
||||
|
||||
"segment_size": 8192,
|
||||
"num_mels": 80,
|
||||
"num_freq": 1025,
|
||||
"n_fft": 1024,
|
||||
"hop_size": 256,
|
||||
"win_size": 1024,
|
||||
|
||||
"sampling_rate": 22050,
|
||||
|
||||
"fmin": 0,
|
||||
"fmax": 8000,
|
||||
"fmax_for_loss": null,
|
||||
|
||||
"num_workers": 4,
|
||||
|
||||
"dist_config": {
|
||||
"dist_backend": "nccl",
|
||||
"dist_url": "tcp://localhost:54321",
|
||||
"world_size": 1
|
||||
}
|
||||
}
|
||||
@@ -1,45 +0,0 @@
|
||||
{
|
||||
"resblock": "1",
|
||||
"num_gpus": 0,
|
||||
"batch_size": 32,
|
||||
"learning_rate": 0.0001,
|
||||
"adam_b1": 0.8,
|
||||
"adam_b2": 0.99,
|
||||
"lr_decay": 0.9999996,
|
||||
"seed": 1234,
|
||||
|
||||
"upsample_rates": [4,4,2,2,2,2],
|
||||
"upsample_kernel_sizes": [8,8,4,4,4,4],
|
||||
"upsample_initial_channel": 1536,
|
||||
"resblock_kernel_sizes": [3,7,11],
|
||||
"resblock_dilation_sizes": [[1,3,5], [1,3,5], [1,3,5]],
|
||||
|
||||
"activation": "snakebeta",
|
||||
"snake_logscale": true,
|
||||
|
||||
"resolutions": [[1024, 120, 600], [2048, 240, 1200], [512, 50, 240]],
|
||||
"mpd_reshapes": [2, 3, 5, 7, 11],
|
||||
"use_spectral_norm": false,
|
||||
"discriminator_channel_mult": 1,
|
||||
|
||||
"segment_size": 8192,
|
||||
"num_mels": 100,
|
||||
"num_freq": 1025,
|
||||
"n_fft": 1024,
|
||||
"hop_size": 256,
|
||||
"win_size": 1024,
|
||||
|
||||
"sampling_rate": 24000,
|
||||
|
||||
"fmin": 0,
|
||||
"fmax": 12000,
|
||||
"fmax_for_loss": null,
|
||||
|
||||
"num_workers": 4,
|
||||
|
||||
"dist_config": {
|
||||
"dist_backend": "nccl",
|
||||
"dist_url": "tcp://localhost:54321",
|
||||
"world_size": 1
|
||||
}
|
||||
}
|
||||
@@ -1,45 +0,0 @@
|
||||
{
|
||||
"resblock": "1",
|
||||
"num_gpus": 0,
|
||||
"batch_size": 32,
|
||||
"learning_rate": 0.0001,
|
||||
"adam_b1": 0.8,
|
||||
"adam_b2": 0.99,
|
||||
"lr_decay": 0.9999996,
|
||||
"seed": 1234,
|
||||
|
||||
"upsample_rates": [8,8,2,2],
|
||||
"upsample_kernel_sizes": [16,16,4,4],
|
||||
"upsample_initial_channel": 512,
|
||||
"resblock_kernel_sizes": [3,7,11],
|
||||
"resblock_dilation_sizes": [[1,3,5], [1,3,5], [1,3,5]],
|
||||
|
||||
"activation": "snakebeta",
|
||||
"snake_logscale": true,
|
||||
|
||||
"resolutions": [[1024, 120, 600], [2048, 240, 1200], [512, 50, 240]],
|
||||
"mpd_reshapes": [2, 3, 5, 7, 11],
|
||||
"use_spectral_norm": false,
|
||||
"discriminator_channel_mult": 1,
|
||||
|
||||
"segment_size": 8192,
|
||||
"num_mels": 80,
|
||||
"num_freq": 1025,
|
||||
"n_fft": 1024,
|
||||
"hop_size": 256,
|
||||
"win_size": 1024,
|
||||
|
||||
"sampling_rate": 22050,
|
||||
|
||||
"fmin": 0,
|
||||
"fmax": 8000,
|
||||
"fmax_for_loss": null,
|
||||
|
||||
"num_workers": 4,
|
||||
|
||||
"dist_config": {
|
||||
"dist_backend": "nccl",
|
||||
"dist_url": "tcp://localhost:54321",
|
||||
"world_size": 1
|
||||
}
|
||||
}
|
||||
@@ -1,45 +0,0 @@
|
||||
{
|
||||
"resblock": "1",
|
||||
"num_gpus": 0,
|
||||
"batch_size": 32,
|
||||
"learning_rate": 0.0001,
|
||||
"adam_b1": 0.8,
|
||||
"adam_b2": 0.99,
|
||||
"lr_decay": 0.9999996,
|
||||
"seed": 1234,
|
||||
|
||||
"upsample_rates": [8,8,2,2],
|
||||
"upsample_kernel_sizes": [16,16,4,4],
|
||||
"upsample_initial_channel": 512,
|
||||
"resblock_kernel_sizes": [3,7,11],
|
||||
"resblock_dilation_sizes": [[1,3,5], [1,3,5], [1,3,5]],
|
||||
|
||||
"activation": "snakebeta",
|
||||
"snake_logscale": true,
|
||||
|
||||
"resolutions": [[1024, 120, 600], [2048, 240, 1200], [512, 50, 240]],
|
||||
"mpd_reshapes": [2, 3, 5, 7, 11],
|
||||
"use_spectral_norm": false,
|
||||
"discriminator_channel_mult": 1,
|
||||
|
||||
"segment_size": 8192,
|
||||
"num_mels": 100,
|
||||
"num_freq": 1025,
|
||||
"n_fft": 1024,
|
||||
"hop_size": 256,
|
||||
"win_size": 1024,
|
||||
|
||||
"sampling_rate": 24000,
|
||||
|
||||
"fmin": 0,
|
||||
"fmax": 12000,
|
||||
"fmax_for_loss": null,
|
||||
|
||||
"num_workers": 4,
|
||||
|
||||
"dist_config": {
|
||||
"dist_backend": "nccl",
|
||||
"dist_url": "tcp://localhost:54321",
|
||||
"world_size": 1
|
||||
}
|
||||
}
|
||||
@@ -1,61 +0,0 @@
|
||||
{
|
||||
"resblock": "1",
|
||||
"num_gpus": 0,
|
||||
"batch_size": 4,
|
||||
"learning_rate": 0.0001,
|
||||
"adam_b1": 0.8,
|
||||
"adam_b2": 0.99,
|
||||
"lr_decay": 0.9999996,
|
||||
"seed": 1234,
|
||||
|
||||
"upsample_rates": [4,4,2,2,2,2],
|
||||
"upsample_kernel_sizes": [8,8,4,4,4,4],
|
||||
"upsample_initial_channel": 1536,
|
||||
"resblock_kernel_sizes": [3,7,11],
|
||||
"resblock_dilation_sizes": [[1,3,5], [1,3,5], [1,3,5]],
|
||||
|
||||
"use_tanh_at_final": false,
|
||||
"use_bias_at_final": false,
|
||||
|
||||
"activation": "snakebeta",
|
||||
"snake_logscale": true,
|
||||
|
||||
"use_cqtd_instead_of_mrd": true,
|
||||
"cqtd_filters": 128,
|
||||
"cqtd_max_filters": 1024,
|
||||
"cqtd_filters_scale": 1,
|
||||
"cqtd_dilations": [1, 2, 4],
|
||||
"cqtd_hop_lengths": [512, 256, 256],
|
||||
"cqtd_n_octaves": [9, 9, 9],
|
||||
"cqtd_bins_per_octaves": [24, 36, 48],
|
||||
|
||||
"mpd_reshapes": [2, 3, 5, 7, 11],
|
||||
"use_spectral_norm": false,
|
||||
"discriminator_channel_mult": 1,
|
||||
|
||||
"use_multiscale_melloss": true,
|
||||
"lambda_melloss": 15,
|
||||
|
||||
"clip_grad_norm": 500,
|
||||
|
||||
"segment_size": 65536,
|
||||
"num_mels": 80,
|
||||
"num_freq": 1025,
|
||||
"n_fft": 1024,
|
||||
"hop_size": 256,
|
||||
"win_size": 1024,
|
||||
|
||||
"sampling_rate": 22050,
|
||||
|
||||
"fmin": 0,
|
||||
"fmax": null,
|
||||
"fmax_for_loss": null,
|
||||
|
||||
"num_workers": 4,
|
||||
|
||||
"dist_config": {
|
||||
"dist_backend": "nccl",
|
||||
"dist_url": "tcp://localhost:54321",
|
||||
"world_size": 1
|
||||
}
|
||||
}
|
||||
@@ -1,61 +0,0 @@
|
||||
{
|
||||
"resblock": "1",
|
||||
"num_gpus": 0,
|
||||
"batch_size": 4,
|
||||
"learning_rate": 0.0001,
|
||||
"adam_b1": 0.8,
|
||||
"adam_b2": 0.99,
|
||||
"lr_decay": 0.9999996,
|
||||
"seed": 1234,
|
||||
|
||||
"upsample_rates": [4,4,2,2,2,2],
|
||||
"upsample_kernel_sizes": [8,8,4,4,4,4],
|
||||
"upsample_initial_channel": 1536,
|
||||
"resblock_kernel_sizes": [3,7,11],
|
||||
"resblock_dilation_sizes": [[1,3,5], [1,3,5], [1,3,5]],
|
||||
|
||||
"use_tanh_at_final": false,
|
||||
"use_bias_at_final": false,
|
||||
|
||||
"activation": "snakebeta",
|
||||
"snake_logscale": true,
|
||||
|
||||
"use_cqtd_instead_of_mrd": true,
|
||||
"cqtd_filters": 128,
|
||||
"cqtd_max_filters": 1024,
|
||||
"cqtd_filters_scale": 1,
|
||||
"cqtd_dilations": [1, 2, 4],
|
||||
"cqtd_hop_lengths": [512, 256, 256],
|
||||
"cqtd_n_octaves": [9, 9, 9],
|
||||
"cqtd_bins_per_octaves": [24, 36, 48],
|
||||
|
||||
"mpd_reshapes": [2, 3, 5, 7, 11],
|
||||
"use_spectral_norm": false,
|
||||
"discriminator_channel_mult": 1,
|
||||
|
||||
"use_multiscale_melloss": true,
|
||||
"lambda_melloss": 15,
|
||||
|
||||
"clip_grad_norm": 500,
|
||||
|
||||
"segment_size": 65536,
|
||||
"num_mels": 80,
|
||||
"num_freq": 1025,
|
||||
"n_fft": 1024,
|
||||
"hop_size": 256,
|
||||
"win_size": 1024,
|
||||
|
||||
"sampling_rate": 22050,
|
||||
|
||||
"fmin": 0,
|
||||
"fmax": 8000,
|
||||
"fmax_for_loss": null,
|
||||
|
||||
"num_workers": 4,
|
||||
|
||||
"dist_config": {
|
||||
"dist_backend": "nccl",
|
||||
"dist_url": "tcp://localhost:54321",
|
||||
"world_size": 1
|
||||
}
|
||||
}
|
||||
@@ -1,61 +0,0 @@
|
||||
{
|
||||
"resblock": "1",
|
||||
"num_gpus": 0,
|
||||
"batch_size": 4,
|
||||
"learning_rate": 0.0001,
|
||||
"adam_b1": 0.8,
|
||||
"adam_b2": 0.99,
|
||||
"lr_decay": 0.9999996,
|
||||
"seed": 1234,
|
||||
|
||||
"upsample_rates": [4,4,2,2,2,2],
|
||||
"upsample_kernel_sizes": [8,8,4,4,4,4],
|
||||
"upsample_initial_channel": 1536,
|
||||
"resblock_kernel_sizes": [3,7,11],
|
||||
"resblock_dilation_sizes": [[1,3,5], [1,3,5], [1,3,5]],
|
||||
|
||||
"use_tanh_at_final": false,
|
||||
"use_bias_at_final": false,
|
||||
|
||||
"activation": "snakebeta",
|
||||
"snake_logscale": true,
|
||||
|
||||
"use_cqtd_instead_of_mrd": true,
|
||||
"cqtd_filters": 128,
|
||||
"cqtd_max_filters": 1024,
|
||||
"cqtd_filters_scale": 1,
|
||||
"cqtd_dilations": [1, 2, 4],
|
||||
"cqtd_hop_lengths": [512, 256, 256],
|
||||
"cqtd_n_octaves": [9, 9, 9],
|
||||
"cqtd_bins_per_octaves": [24, 36, 48],
|
||||
|
||||
"mpd_reshapes": [2, 3, 5, 7, 11],
|
||||
"use_spectral_norm": false,
|
||||
"discriminator_channel_mult": 1,
|
||||
|
||||
"use_multiscale_melloss": true,
|
||||
"lambda_melloss": 15,
|
||||
|
||||
"clip_grad_norm": 500,
|
||||
|
||||
"segment_size": 65536,
|
||||
"num_mels": 100,
|
||||
"num_freq": 1025,
|
||||
"n_fft": 1024,
|
||||
"hop_size": 256,
|
||||
"win_size": 1024,
|
||||
|
||||
"sampling_rate": 24000,
|
||||
|
||||
"fmin": 0,
|
||||
"fmax": null,
|
||||
"fmax_for_loss": null,
|
||||
|
||||
"num_workers": 4,
|
||||
|
||||
"dist_config": {
|
||||
"dist_backend": "nccl",
|
||||
"dist_url": "tcp://localhost:54321",
|
||||
"world_size": 1
|
||||
}
|
||||
}
|
||||
@@ -1,61 +0,0 @@
|
||||
{
|
||||
"resblock": "1",
|
||||
"num_gpus": 0,
|
||||
"batch_size": 4,
|
||||
"learning_rate": 0.0001,
|
||||
"adam_b1": 0.8,
|
||||
"adam_b2": 0.99,
|
||||
"lr_decay": 0.9999996,
|
||||
"seed": 1234,
|
||||
|
||||
"upsample_rates": [4,4,2,2,2,2],
|
||||
"upsample_kernel_sizes": [8,8,4,4,4,4],
|
||||
"upsample_initial_channel": 1536,
|
||||
"resblock_kernel_sizes": [3,7,11],
|
||||
"resblock_dilation_sizes": [[1,3,5], [1,3,5], [1,3,5]],
|
||||
|
||||
"use_tanh_at_final": false,
|
||||
"use_bias_at_final": false,
|
||||
|
||||
"activation": "snakebeta",
|
||||
"snake_logscale": true,
|
||||
|
||||
"use_cqtd_instead_of_mrd": true,
|
||||
"cqtd_filters": 128,
|
||||
"cqtd_max_filters": 1024,
|
||||
"cqtd_filters_scale": 1,
|
||||
"cqtd_dilations": [1, 2, 4],
|
||||
"cqtd_hop_lengths": [512, 256, 256],
|
||||
"cqtd_n_octaves": [9, 9, 9],
|
||||
"cqtd_bins_per_octaves": [24, 36, 48],
|
||||
|
||||
"mpd_reshapes": [2, 3, 5, 7, 11],
|
||||
"use_spectral_norm": false,
|
||||
"discriminator_channel_mult": 1,
|
||||
|
||||
"use_multiscale_melloss": true,
|
||||
"lambda_melloss": 15,
|
||||
|
||||
"clip_grad_norm": 500,
|
||||
|
||||
"segment_size": 65536,
|
||||
"num_mels": 128,
|
||||
"num_freq": 1025,
|
||||
"n_fft": 1024,
|
||||
"hop_size": 256,
|
||||
"win_size": 1024,
|
||||
|
||||
"sampling_rate": 44100,
|
||||
|
||||
"fmin": 0,
|
||||
"fmax": null,
|
||||
"fmax_for_loss": null,
|
||||
|
||||
"num_workers": 4,
|
||||
|
||||
"dist_config": {
|
||||
"dist_backend": "nccl",
|
||||
"dist_url": "tcp://localhost:54321",
|
||||
"world_size": 1
|
||||
}
|
||||
}
|
||||
@@ -1,61 +0,0 @@
|
||||
{
|
||||
"resblock": "1",
|
||||
"num_gpus": 0,
|
||||
"batch_size": 4,
|
||||
"learning_rate": 0.0001,
|
||||
"adam_b1": 0.8,
|
||||
"adam_b2": 0.99,
|
||||
"lr_decay": 0.9999996,
|
||||
"seed": 1234,
|
||||
|
||||
"upsample_rates": [8,4,2,2,2,2],
|
||||
"upsample_kernel_sizes": [16,8,4,4,4,4],
|
||||
"upsample_initial_channel": 1536,
|
||||
"resblock_kernel_sizes": [3,7,11],
|
||||
"resblock_dilation_sizes": [[1,3,5], [1,3,5], [1,3,5]],
|
||||
|
||||
"use_tanh_at_final": false,
|
||||
"use_bias_at_final": false,
|
||||
|
||||
"activation": "snakebeta",
|
||||
"snake_logscale": true,
|
||||
|
||||
"use_cqtd_instead_of_mrd": true,
|
||||
"cqtd_filters": 128,
|
||||
"cqtd_max_filters": 1024,
|
||||
"cqtd_filters_scale": 1,
|
||||
"cqtd_dilations": [1, 2, 4],
|
||||
"cqtd_hop_lengths": [512, 256, 256],
|
||||
"cqtd_n_octaves": [9, 9, 9],
|
||||
"cqtd_bins_per_octaves": [24, 36, 48],
|
||||
|
||||
"mpd_reshapes": [2, 3, 5, 7, 11],
|
||||
"use_spectral_norm": false,
|
||||
"discriminator_channel_mult": 1,
|
||||
|
||||
"use_multiscale_melloss": true,
|
||||
"lambda_melloss": 15,
|
||||
|
||||
"clip_grad_norm": 500,
|
||||
|
||||
"segment_size": 65536,
|
||||
"num_mels": 128,
|
||||
"num_freq": 2049,
|
||||
"n_fft": 2048,
|
||||
"hop_size": 512,
|
||||
"win_size": 2048,
|
||||
|
||||
"sampling_rate": 44100,
|
||||
|
||||
"fmin": 0,
|
||||
"fmax": null,
|
||||
"fmax_for_loss": null,
|
||||
|
||||
"num_workers": 4,
|
||||
|
||||
"dist_config": {
|
||||
"dist_backend": "nccl",
|
||||
"dist_url": "tcp://localhost:54321",
|
||||
"world_size": 1
|
||||
}
|
||||
}
|
||||
@@ -1,625 +0,0 @@
|
||||
# Copyright (c) 2024 NVIDIA CORPORATION.
|
||||
# Licensed under the MIT license.
|
||||
|
||||
# Adapted from https://github.com/jik876/hifi-gan under the MIT license.
|
||||
# LICENSE is in incl_licenses directory.
|
||||
|
||||
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
import torch.nn as nn
|
||||
from torch.nn import Conv2d
|
||||
from torch.nn.utils import weight_norm, spectral_norm
|
||||
from torchaudio.transforms import Spectrogram, Resample
|
||||
|
||||
from env import AttrDict
|
||||
from utils import get_padding
|
||||
import typing
|
||||
from typing import List, Tuple
|
||||
|
||||
|
||||
class DiscriminatorP(torch.nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
h: AttrDict,
|
||||
period: List[int],
|
||||
kernel_size: int = 5,
|
||||
stride: int = 3,
|
||||
use_spectral_norm: bool = False,
|
||||
):
|
||||
super().__init__()
|
||||
self.period = period
|
||||
self.d_mult = h.discriminator_channel_mult
|
||||
norm_f = weight_norm if not use_spectral_norm else spectral_norm
|
||||
|
||||
self.convs = nn.ModuleList(
|
||||
[
|
||||
norm_f(
|
||||
Conv2d(
|
||||
1,
|
||||
int(32 * self.d_mult),
|
||||
(kernel_size, 1),
|
||||
(stride, 1),
|
||||
padding=(get_padding(5, 1), 0),
|
||||
)
|
||||
),
|
||||
norm_f(
|
||||
Conv2d(
|
||||
int(32 * self.d_mult),
|
||||
int(128 * self.d_mult),
|
||||
(kernel_size, 1),
|
||||
(stride, 1),
|
||||
padding=(get_padding(5, 1), 0),
|
||||
)
|
||||
),
|
||||
norm_f(
|
||||
Conv2d(
|
||||
int(128 * self.d_mult),
|
||||
int(512 * self.d_mult),
|
||||
(kernel_size, 1),
|
||||
(stride, 1),
|
||||
padding=(get_padding(5, 1), 0),
|
||||
)
|
||||
),
|
||||
norm_f(
|
||||
Conv2d(
|
||||
int(512 * self.d_mult),
|
||||
int(1024 * self.d_mult),
|
||||
(kernel_size, 1),
|
||||
(stride, 1),
|
||||
padding=(get_padding(5, 1), 0),
|
||||
)
|
||||
),
|
||||
norm_f(
|
||||
Conv2d(
|
||||
int(1024 * self.d_mult),
|
||||
int(1024 * self.d_mult),
|
||||
(kernel_size, 1),
|
||||
1,
|
||||
padding=(2, 0),
|
||||
)
|
||||
),
|
||||
]
|
||||
)
|
||||
self.conv_post = norm_f(Conv2d(int(1024 * self.d_mult), 1, (3, 1), 1, padding=(1, 0)))
|
||||
|
||||
def forward(self, x: torch.Tensor) -> Tuple[torch.Tensor, List[torch.Tensor]]:
|
||||
fmap = []
|
||||
|
||||
# 1d to 2d
|
||||
b, c, t = x.shape
|
||||
if t % self.period != 0: # pad first
|
||||
n_pad = self.period - (t % self.period)
|
||||
x = F.pad(x, (0, n_pad), "reflect")
|
||||
t = t + n_pad
|
||||
x = x.view(b, c, t // self.period, self.period)
|
||||
|
||||
for l in self.convs:
|
||||
x = l(x)
|
||||
x = F.leaky_relu(x, 0.1)
|
||||
fmap.append(x)
|
||||
x = self.conv_post(x)
|
||||
fmap.append(x)
|
||||
x = torch.flatten(x, 1, -1)
|
||||
|
||||
return x, fmap
|
||||
|
||||
|
||||
class MultiPeriodDiscriminator(torch.nn.Module):
|
||||
def __init__(self, h: AttrDict):
|
||||
super().__init__()
|
||||
self.mpd_reshapes = h.mpd_reshapes
|
||||
print(f"mpd_reshapes: {self.mpd_reshapes}")
|
||||
self.discriminators = nn.ModuleList(
|
||||
[DiscriminatorP(h, rs, use_spectral_norm=h.use_spectral_norm) for rs in self.mpd_reshapes]
|
||||
)
|
||||
|
||||
def forward(
|
||||
self, y: torch.Tensor, y_hat: torch.Tensor
|
||||
) -> Tuple[
|
||||
List[torch.Tensor],
|
||||
List[torch.Tensor],
|
||||
List[List[torch.Tensor]],
|
||||
List[List[torch.Tensor]],
|
||||
]:
|
||||
y_d_rs = []
|
||||
y_d_gs = []
|
||||
fmap_rs = []
|
||||
fmap_gs = []
|
||||
for i, d in enumerate(self.discriminators):
|
||||
y_d_r, fmap_r = d(y)
|
||||
y_d_g, fmap_g = d(y_hat)
|
||||
y_d_rs.append(y_d_r)
|
||||
fmap_rs.append(fmap_r)
|
||||
y_d_gs.append(y_d_g)
|
||||
fmap_gs.append(fmap_g)
|
||||
|
||||
return y_d_rs, y_d_gs, fmap_rs, fmap_gs
|
||||
|
||||
|
||||
class DiscriminatorR(nn.Module):
|
||||
def __init__(self, cfg: AttrDict, resolution: List[List[int]]):
|
||||
super().__init__()
|
||||
|
||||
self.resolution = resolution
|
||||
assert len(self.resolution) == 3, f"MRD layer requires list with len=3, got {self.resolution}"
|
||||
self.lrelu_slope = 0.1
|
||||
|
||||
norm_f = weight_norm if cfg.use_spectral_norm == False else spectral_norm
|
||||
if hasattr(cfg, "mrd_use_spectral_norm"):
|
||||
print(f"[INFO] overriding MRD use_spectral_norm as {cfg.mrd_use_spectral_norm}")
|
||||
norm_f = weight_norm if cfg.mrd_use_spectral_norm == False else spectral_norm
|
||||
self.d_mult = cfg.discriminator_channel_mult
|
||||
if hasattr(cfg, "mrd_channel_mult"):
|
||||
print(f"[INFO] overriding mrd channel multiplier as {cfg.mrd_channel_mult}")
|
||||
self.d_mult = cfg.mrd_channel_mult
|
||||
|
||||
self.convs = nn.ModuleList(
|
||||
[
|
||||
norm_f(nn.Conv2d(1, int(32 * self.d_mult), (3, 9), padding=(1, 4))),
|
||||
norm_f(
|
||||
nn.Conv2d(
|
||||
int(32 * self.d_mult),
|
||||
int(32 * self.d_mult),
|
||||
(3, 9),
|
||||
stride=(1, 2),
|
||||
padding=(1, 4),
|
||||
)
|
||||
),
|
||||
norm_f(
|
||||
nn.Conv2d(
|
||||
int(32 * self.d_mult),
|
||||
int(32 * self.d_mult),
|
||||
(3, 9),
|
||||
stride=(1, 2),
|
||||
padding=(1, 4),
|
||||
)
|
||||
),
|
||||
norm_f(
|
||||
nn.Conv2d(
|
||||
int(32 * self.d_mult),
|
||||
int(32 * self.d_mult),
|
||||
(3, 9),
|
||||
stride=(1, 2),
|
||||
padding=(1, 4),
|
||||
)
|
||||
),
|
||||
norm_f(
|
||||
nn.Conv2d(
|
||||
int(32 * self.d_mult),
|
||||
int(32 * self.d_mult),
|
||||
(3, 3),
|
||||
padding=(1, 1),
|
||||
)
|
||||
),
|
||||
]
|
||||
)
|
||||
self.conv_post = norm_f(nn.Conv2d(int(32 * self.d_mult), 1, (3, 3), padding=(1, 1)))
|
||||
|
||||
def forward(self, x: torch.Tensor) -> Tuple[torch.Tensor, List[torch.Tensor]]:
|
||||
fmap = []
|
||||
|
||||
x = self.spectrogram(x)
|
||||
x = x.unsqueeze(1)
|
||||
for l in self.convs:
|
||||
x = l(x)
|
||||
x = F.leaky_relu(x, self.lrelu_slope)
|
||||
fmap.append(x)
|
||||
x = self.conv_post(x)
|
||||
fmap.append(x)
|
||||
x = torch.flatten(x, 1, -1)
|
||||
|
||||
return x, fmap
|
||||
|
||||
def spectrogram(self, x: torch.Tensor) -> torch.Tensor:
|
||||
n_fft, hop_length, win_length = self.resolution
|
||||
x = F.pad(
|
||||
x,
|
||||
(int((n_fft - hop_length) / 2), int((n_fft - hop_length) / 2)),
|
||||
mode="reflect",
|
||||
)
|
||||
x = x.squeeze(1)
|
||||
x = torch.stft(
|
||||
x,
|
||||
n_fft=n_fft,
|
||||
hop_length=hop_length,
|
||||
win_length=win_length,
|
||||
center=False,
|
||||
return_complex=True,
|
||||
)
|
||||
x = torch.view_as_real(x) # [B, F, TT, 2]
|
||||
mag = torch.norm(x, p=2, dim=-1) # [B, F, TT]
|
||||
|
||||
return mag
|
||||
|
||||
|
||||
class MultiResolutionDiscriminator(nn.Module):
|
||||
def __init__(self, cfg, debug=False):
|
||||
super().__init__()
|
||||
self.resolutions = cfg.resolutions
|
||||
assert len(self.resolutions) == 3, (
|
||||
f"MRD requires list of list with len=3, each element having a list with len=3. Got {self.resolutions}"
|
||||
)
|
||||
self.discriminators = nn.ModuleList([DiscriminatorR(cfg, resolution) for resolution in self.resolutions])
|
||||
|
||||
def forward(
|
||||
self, y: torch.Tensor, y_hat: torch.Tensor
|
||||
) -> Tuple[
|
||||
List[torch.Tensor],
|
||||
List[torch.Tensor],
|
||||
List[List[torch.Tensor]],
|
||||
List[List[torch.Tensor]],
|
||||
]:
|
||||
y_d_rs = []
|
||||
y_d_gs = []
|
||||
fmap_rs = []
|
||||
fmap_gs = []
|
||||
|
||||
for i, d in enumerate(self.discriminators):
|
||||
y_d_r, fmap_r = d(x=y)
|
||||
y_d_g, fmap_g = d(x=y_hat)
|
||||
y_d_rs.append(y_d_r)
|
||||
fmap_rs.append(fmap_r)
|
||||
y_d_gs.append(y_d_g)
|
||||
fmap_gs.append(fmap_g)
|
||||
|
||||
return y_d_rs, y_d_gs, fmap_rs, fmap_gs
|
||||
|
||||
|
||||
# Method based on descript-audio-codec: https://github.com/descriptinc/descript-audio-codec
|
||||
# Modified code adapted from https://github.com/gemelo-ai/vocos under the MIT license.
|
||||
# LICENSE is in incl_licenses directory.
|
||||
class DiscriminatorB(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
window_length: int,
|
||||
channels: int = 32,
|
||||
hop_factor: float = 0.25,
|
||||
bands: Tuple[Tuple[float, float], ...] = (
|
||||
(0.0, 0.1),
|
||||
(0.1, 0.25),
|
||||
(0.25, 0.5),
|
||||
(0.5, 0.75),
|
||||
(0.75, 1.0),
|
||||
),
|
||||
):
|
||||
super().__init__()
|
||||
self.window_length = window_length
|
||||
self.hop_factor = hop_factor
|
||||
self.spec_fn = Spectrogram(
|
||||
n_fft=window_length,
|
||||
hop_length=int(window_length * hop_factor),
|
||||
win_length=window_length,
|
||||
power=None,
|
||||
)
|
||||
n_fft = window_length // 2 + 1
|
||||
bands = [(int(b[0] * n_fft), int(b[1] * n_fft)) for b in bands]
|
||||
self.bands = bands
|
||||
convs = lambda: nn.ModuleList(
|
||||
[
|
||||
weight_norm(nn.Conv2d(2, channels, (3, 9), (1, 1), padding=(1, 4))),
|
||||
weight_norm(nn.Conv2d(channels, channels, (3, 9), (1, 2), padding=(1, 4))),
|
||||
weight_norm(nn.Conv2d(channels, channels, (3, 9), (1, 2), padding=(1, 4))),
|
||||
weight_norm(nn.Conv2d(channels, channels, (3, 9), (1, 2), padding=(1, 4))),
|
||||
weight_norm(nn.Conv2d(channels, channels, (3, 3), (1, 1), padding=(1, 1))),
|
||||
]
|
||||
)
|
||||
self.band_convs = nn.ModuleList([convs() for _ in range(len(self.bands))])
|
||||
|
||||
self.conv_post = weight_norm(nn.Conv2d(channels, 1, (3, 3), (1, 1), padding=(1, 1)))
|
||||
|
||||
def spectrogram(self, x: torch.Tensor) -> List[torch.Tensor]:
|
||||
# Remove DC offset
|
||||
x = x - x.mean(dim=-1, keepdims=True)
|
||||
# Peak normalize the volume of input audio
|
||||
x = 0.8 * x / (x.abs().max(dim=-1, keepdim=True)[0] + 1e-9)
|
||||
x = self.spec_fn(x)
|
||||
x = torch.view_as_real(x)
|
||||
x = x.permute(0, 3, 2, 1) # [B, F, T, C] -> [B, C, T, F]
|
||||
# Split into bands
|
||||
x_bands = [x[..., b[0] : b[1]] for b in self.bands]
|
||||
return x_bands
|
||||
|
||||
def forward(self, x: torch.Tensor) -> Tuple[torch.Tensor, List[torch.Tensor]]:
|
||||
x_bands = self.spectrogram(x.squeeze(1))
|
||||
fmap = []
|
||||
x = []
|
||||
|
||||
for band, stack in zip(x_bands, self.band_convs):
|
||||
for i, layer in enumerate(stack):
|
||||
band = layer(band)
|
||||
band = torch.nn.functional.leaky_relu(band, 0.1)
|
||||
if i > 0:
|
||||
fmap.append(band)
|
||||
x.append(band)
|
||||
|
||||
x = torch.cat(x, dim=-1)
|
||||
x = self.conv_post(x)
|
||||
fmap.append(x)
|
||||
|
||||
return x, fmap
|
||||
|
||||
|
||||
# Method based on descript-audio-codec: https://github.com/descriptinc/descript-audio-codec
|
||||
# Modified code adapted from https://github.com/gemelo-ai/vocos under the MIT license.
|
||||
# LICENSE is in incl_licenses directory.
|
||||
class MultiBandDiscriminator(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
h,
|
||||
):
|
||||
"""
|
||||
Multi-band multi-scale STFT discriminator, with the architecture based on https://github.com/descriptinc/descript-audio-codec.
|
||||
and the modified code adapted from https://github.com/gemelo-ai/vocos.
|
||||
"""
|
||||
super().__init__()
|
||||
# fft_sizes (list[int]): Tuple of window lengths for FFT. Defaults to [2048, 1024, 512] if not set in h.
|
||||
self.fft_sizes = h.get("mbd_fft_sizes", [2048, 1024, 512])
|
||||
self.discriminators = nn.ModuleList([DiscriminatorB(window_length=w) for w in self.fft_sizes])
|
||||
|
||||
def forward(
|
||||
self, y: torch.Tensor, y_hat: torch.Tensor
|
||||
) -> Tuple[
|
||||
List[torch.Tensor],
|
||||
List[torch.Tensor],
|
||||
List[List[torch.Tensor]],
|
||||
List[List[torch.Tensor]],
|
||||
]:
|
||||
y_d_rs = []
|
||||
y_d_gs = []
|
||||
fmap_rs = []
|
||||
fmap_gs = []
|
||||
|
||||
for d in self.discriminators:
|
||||
y_d_r, fmap_r = d(x=y)
|
||||
y_d_g, fmap_g = d(x=y_hat)
|
||||
y_d_rs.append(y_d_r)
|
||||
fmap_rs.append(fmap_r)
|
||||
y_d_gs.append(y_d_g)
|
||||
fmap_gs.append(fmap_g)
|
||||
|
||||
return y_d_rs, y_d_gs, fmap_rs, fmap_gs
|
||||
|
||||
|
||||
# Adapted from https://github.com/open-mmlab/Amphion/blob/main/models/vocoders/gan/discriminator/mssbcqtd.py under the MIT license.
|
||||
# LICENSE is in incl_licenses directory.
|
||||
class DiscriminatorCQT(nn.Module):
|
||||
def __init__(self, cfg: AttrDict, hop_length: int, n_octaves: int, bins_per_octave: int):
|
||||
super().__init__()
|
||||
self.cfg = cfg
|
||||
|
||||
self.filters = cfg["cqtd_filters"]
|
||||
self.max_filters = cfg["cqtd_max_filters"]
|
||||
self.filters_scale = cfg["cqtd_filters_scale"]
|
||||
self.kernel_size = (3, 9)
|
||||
self.dilations = cfg["cqtd_dilations"]
|
||||
self.stride = (1, 2)
|
||||
|
||||
self.in_channels = cfg["cqtd_in_channels"]
|
||||
self.out_channels = cfg["cqtd_out_channels"]
|
||||
self.fs = cfg["sampling_rate"]
|
||||
self.hop_length = hop_length
|
||||
self.n_octaves = n_octaves
|
||||
self.bins_per_octave = bins_per_octave
|
||||
|
||||
# Lazy-load
|
||||
from nnAudio import features
|
||||
|
||||
self.cqt_transform = features.cqt.CQT2010v2(
|
||||
sr=self.fs * 2,
|
||||
hop_length=self.hop_length,
|
||||
n_bins=self.bins_per_octave * self.n_octaves,
|
||||
bins_per_octave=self.bins_per_octave,
|
||||
output_format="Complex",
|
||||
pad_mode="constant",
|
||||
)
|
||||
|
||||
self.conv_pres = nn.ModuleList()
|
||||
for _ in range(self.n_octaves):
|
||||
self.conv_pres.append(
|
||||
nn.Conv2d(
|
||||
self.in_channels * 2,
|
||||
self.in_channels * 2,
|
||||
kernel_size=self.kernel_size,
|
||||
padding=self.get_2d_padding(self.kernel_size),
|
||||
)
|
||||
)
|
||||
|
||||
self.convs = nn.ModuleList()
|
||||
|
||||
self.convs.append(
|
||||
nn.Conv2d(
|
||||
self.in_channels * 2,
|
||||
self.filters,
|
||||
kernel_size=self.kernel_size,
|
||||
padding=self.get_2d_padding(self.kernel_size),
|
||||
)
|
||||
)
|
||||
|
||||
in_chs = min(self.filters_scale * self.filters, self.max_filters)
|
||||
for i, dilation in enumerate(self.dilations):
|
||||
out_chs = min((self.filters_scale ** (i + 1)) * self.filters, self.max_filters)
|
||||
self.convs.append(
|
||||
weight_norm(
|
||||
nn.Conv2d(
|
||||
in_chs,
|
||||
out_chs,
|
||||
kernel_size=self.kernel_size,
|
||||
stride=self.stride,
|
||||
dilation=(dilation, 1),
|
||||
padding=self.get_2d_padding(self.kernel_size, (dilation, 1)),
|
||||
)
|
||||
)
|
||||
)
|
||||
in_chs = out_chs
|
||||
out_chs = min(
|
||||
(self.filters_scale ** (len(self.dilations) + 1)) * self.filters,
|
||||
self.max_filters,
|
||||
)
|
||||
self.convs.append(
|
||||
weight_norm(
|
||||
nn.Conv2d(
|
||||
in_chs,
|
||||
out_chs,
|
||||
kernel_size=(self.kernel_size[0], self.kernel_size[0]),
|
||||
padding=self.get_2d_padding((self.kernel_size[0], self.kernel_size[0])),
|
||||
)
|
||||
)
|
||||
)
|
||||
|
||||
self.conv_post = weight_norm(
|
||||
nn.Conv2d(
|
||||
out_chs,
|
||||
self.out_channels,
|
||||
kernel_size=(self.kernel_size[0], self.kernel_size[0]),
|
||||
padding=self.get_2d_padding((self.kernel_size[0], self.kernel_size[0])),
|
||||
)
|
||||
)
|
||||
|
||||
self.activation = torch.nn.LeakyReLU(negative_slope=0.1)
|
||||
self.resample = Resample(orig_freq=self.fs, new_freq=self.fs * 2)
|
||||
|
||||
self.cqtd_normalize_volume = self.cfg.get("cqtd_normalize_volume", False)
|
||||
if self.cqtd_normalize_volume:
|
||||
print(
|
||||
"[INFO] cqtd_normalize_volume set to True. Will apply DC offset removal & peak volume normalization in CQTD!"
|
||||
)
|
||||
|
||||
def get_2d_padding(
|
||||
self,
|
||||
kernel_size: typing.Tuple[int, int],
|
||||
dilation: typing.Tuple[int, int] = (1, 1),
|
||||
):
|
||||
return (
|
||||
((kernel_size[0] - 1) * dilation[0]) // 2,
|
||||
((kernel_size[1] - 1) * dilation[1]) // 2,
|
||||
)
|
||||
|
||||
def forward(self, x: torch.tensor) -> Tuple[torch.Tensor, List[torch.Tensor]]:
|
||||
fmap = []
|
||||
|
||||
if self.cqtd_normalize_volume:
|
||||
# Remove DC offset
|
||||
x = x - x.mean(dim=-1, keepdims=True)
|
||||
# Peak normalize the volume of input audio
|
||||
x = 0.8 * x / (x.abs().max(dim=-1, keepdim=True)[0] + 1e-9)
|
||||
|
||||
x = self.resample(x)
|
||||
|
||||
z = self.cqt_transform(x)
|
||||
|
||||
z_amplitude = z[:, :, :, 0].unsqueeze(1)
|
||||
z_phase = z[:, :, :, 1].unsqueeze(1)
|
||||
|
||||
z = torch.cat([z_amplitude, z_phase], dim=1)
|
||||
z = torch.permute(z, (0, 1, 3, 2)) # [B, C, W, T] -> [B, C, T, W]
|
||||
|
||||
latent_z = []
|
||||
for i in range(self.n_octaves):
|
||||
latent_z.append(
|
||||
self.conv_pres[i](
|
||||
z[
|
||||
:,
|
||||
:,
|
||||
:,
|
||||
i * self.bins_per_octave : (i + 1) * self.bins_per_octave,
|
||||
]
|
||||
)
|
||||
)
|
||||
latent_z = torch.cat(latent_z, dim=-1)
|
||||
|
||||
for i, l in enumerate(self.convs):
|
||||
latent_z = l(latent_z)
|
||||
|
||||
latent_z = self.activation(latent_z)
|
||||
fmap.append(latent_z)
|
||||
|
||||
latent_z = self.conv_post(latent_z)
|
||||
|
||||
return latent_z, fmap
|
||||
|
||||
|
||||
class MultiScaleSubbandCQTDiscriminator(nn.Module):
|
||||
def __init__(self, cfg: AttrDict):
|
||||
super().__init__()
|
||||
|
||||
self.cfg = cfg
|
||||
# Using get with defaults
|
||||
self.cfg["cqtd_filters"] = self.cfg.get("cqtd_filters", 32)
|
||||
self.cfg["cqtd_max_filters"] = self.cfg.get("cqtd_max_filters", 1024)
|
||||
self.cfg["cqtd_filters_scale"] = self.cfg.get("cqtd_filters_scale", 1)
|
||||
self.cfg["cqtd_dilations"] = self.cfg.get("cqtd_dilations", [1, 2, 4])
|
||||
self.cfg["cqtd_in_channels"] = self.cfg.get("cqtd_in_channels", 1)
|
||||
self.cfg["cqtd_out_channels"] = self.cfg.get("cqtd_out_channels", 1)
|
||||
# Multi-scale params to loop over
|
||||
self.cfg["cqtd_hop_lengths"] = self.cfg.get("cqtd_hop_lengths", [512, 256, 256])
|
||||
self.cfg["cqtd_n_octaves"] = self.cfg.get("cqtd_n_octaves", [9, 9, 9])
|
||||
self.cfg["cqtd_bins_per_octaves"] = self.cfg.get("cqtd_bins_per_octaves", [24, 36, 48])
|
||||
|
||||
self.discriminators = nn.ModuleList(
|
||||
[
|
||||
DiscriminatorCQT(
|
||||
self.cfg,
|
||||
hop_length=self.cfg["cqtd_hop_lengths"][i],
|
||||
n_octaves=self.cfg["cqtd_n_octaves"][i],
|
||||
bins_per_octave=self.cfg["cqtd_bins_per_octaves"][i],
|
||||
)
|
||||
for i in range(len(self.cfg["cqtd_hop_lengths"]))
|
||||
]
|
||||
)
|
||||
|
||||
def forward(
|
||||
self, y: torch.Tensor, y_hat: torch.Tensor
|
||||
) -> Tuple[
|
||||
List[torch.Tensor],
|
||||
List[torch.Tensor],
|
||||
List[List[torch.Tensor]],
|
||||
List[List[torch.Tensor]],
|
||||
]:
|
||||
y_d_rs = []
|
||||
y_d_gs = []
|
||||
fmap_rs = []
|
||||
fmap_gs = []
|
||||
|
||||
for disc in self.discriminators:
|
||||
y_d_r, fmap_r = disc(y)
|
||||
y_d_g, fmap_g = disc(y_hat)
|
||||
y_d_rs.append(y_d_r)
|
||||
fmap_rs.append(fmap_r)
|
||||
y_d_gs.append(y_d_g)
|
||||
fmap_gs.append(fmap_g)
|
||||
|
||||
return y_d_rs, y_d_gs, fmap_rs, fmap_gs
|
||||
|
||||
|
||||
class CombinedDiscriminator(nn.Module):
|
||||
"""
|
||||
Wrapper of chaining multiple discrimiantor architectures.
|
||||
Example: combine mbd and cqtd as a single class
|
||||
"""
|
||||
|
||||
def __init__(self, list_discriminator: List[nn.Module]):
|
||||
super().__init__()
|
||||
self.discrimiantor = nn.ModuleList(list_discriminator)
|
||||
|
||||
def forward(
|
||||
self, y: torch.Tensor, y_hat: torch.Tensor
|
||||
) -> Tuple[
|
||||
List[torch.Tensor],
|
||||
List[torch.Tensor],
|
||||
List[List[torch.Tensor]],
|
||||
List[List[torch.Tensor]],
|
||||
]:
|
||||
y_d_rs = []
|
||||
y_d_gs = []
|
||||
fmap_rs = []
|
||||
fmap_gs = []
|
||||
|
||||
for disc in self.discrimiantor:
|
||||
y_d_r, y_d_g, fmap_r, fmap_g = disc(y, y_hat)
|
||||
y_d_rs.extend(y_d_r)
|
||||
fmap_rs.extend(fmap_r)
|
||||
y_d_gs.extend(y_d_g)
|
||||
fmap_gs.extend(fmap_g)
|
||||
|
||||
return y_d_rs, y_d_gs, fmap_rs, fmap_gs
|
||||
@@ -1,18 +1,8 @@
|
||||
# Adapted from https://github.com/jik876/hifi-gan under the MIT license.
|
||||
# LICENSE is in incl_licenses directory.
|
||||
|
||||
import os
|
||||
import shutil
|
||||
# LICENSE is in the incl_licenses directory.
|
||||
|
||||
|
||||
class AttrDict(dict):
|
||||
def __init__(self, *args, **kwargs):
|
||||
super(AttrDict, self).__init__(*args, **kwargs)
|
||||
self.__dict__ = self
|
||||
|
||||
|
||||
def build_env(config, config_name, path):
|
||||
t_path = os.path.join(path, config_name)
|
||||
if config != t_path:
|
||||
os.makedirs(path, exist_ok=True)
|
||||
shutil.copyfile(config, os.path.join(path, config_name))
|
||||
|
||||
@@ -1,85 +0,0 @@
|
||||
# Adapted from https://github.com/jik876/hifi-gan under the MIT license.
|
||||
# LICENSE is in incl_licenses directory.
|
||||
|
||||
from __future__ import absolute_import, division, print_function, unicode_literals
|
||||
|
||||
import os
|
||||
import argparse
|
||||
import json
|
||||
import torch
|
||||
import librosa
|
||||
from utils import load_checkpoint
|
||||
from meldataset import get_mel_spectrogram
|
||||
from scipy.io.wavfile import write
|
||||
from env import AttrDict
|
||||
from meldataset import MAX_WAV_VALUE
|
||||
from bigvgan import BigVGAN as Generator
|
||||
|
||||
h = None
|
||||
device = None
|
||||
torch.backends.cudnn.benchmark = False
|
||||
|
||||
|
||||
def inference(a, h):
|
||||
generator = Generator(h, use_cuda_kernel=a.use_cuda_kernel).to(device)
|
||||
|
||||
state_dict_g = load_checkpoint(a.checkpoint_file, device)
|
||||
generator.load_state_dict(state_dict_g["generator"])
|
||||
|
||||
filelist = os.listdir(a.input_wavs_dir)
|
||||
|
||||
os.makedirs(a.output_dir, exist_ok=True)
|
||||
|
||||
generator.eval()
|
||||
generator.remove_weight_norm()
|
||||
with torch.no_grad():
|
||||
for i, filname in enumerate(filelist):
|
||||
# Load the ground truth audio and resample if necessary
|
||||
wav, sr = librosa.load(os.path.join(a.input_wavs_dir, filname), sr=h.sampling_rate, mono=True)
|
||||
wav = torch.FloatTensor(wav).to(device)
|
||||
# Compute mel spectrogram from the ground truth audio
|
||||
x = get_mel_spectrogram(wav.unsqueeze(0), generator.h)
|
||||
|
||||
y_g_hat = generator(x)
|
||||
|
||||
audio = y_g_hat.squeeze()
|
||||
audio = audio * MAX_WAV_VALUE
|
||||
audio = audio.cpu().numpy().astype("int16")
|
||||
|
||||
output_file = os.path.join(a.output_dir, os.path.splitext(filname)[0] + "_generated.wav")
|
||||
write(output_file, h.sampling_rate, audio)
|
||||
print(output_file)
|
||||
|
||||
|
||||
def main():
|
||||
print("Initializing Inference Process..")
|
||||
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument("--input_wavs_dir", default="test_files")
|
||||
parser.add_argument("--output_dir", default="generated_files")
|
||||
parser.add_argument("--checkpoint_file", required=True)
|
||||
parser.add_argument("--use_cuda_kernel", action="store_true", default=False)
|
||||
|
||||
a = parser.parse_args()
|
||||
|
||||
config_file = os.path.join(os.path.split(a.checkpoint_file)[0], "config.json")
|
||||
with open(config_file) as f:
|
||||
data = f.read()
|
||||
|
||||
global h
|
||||
json_config = json.loads(data)
|
||||
h = AttrDict(json_config)
|
||||
|
||||
torch.manual_seed(h.seed)
|
||||
global device
|
||||
if torch.cuda.is_available():
|
||||
torch.cuda.manual_seed(h.seed)
|
||||
device = torch.device("cuda")
|
||||
else:
|
||||
device = torch.device("cpu")
|
||||
|
||||
inference(a, h)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -1,100 +0,0 @@
|
||||
# Adapted from https://github.com/jik876/hifi-gan under the MIT license.
|
||||
# LICENSE is in incl_licenses directory.
|
||||
|
||||
from __future__ import absolute_import, division, print_function, unicode_literals
|
||||
|
||||
import glob
|
||||
import os
|
||||
import numpy as np
|
||||
import argparse
|
||||
import json
|
||||
import torch
|
||||
from scipy.io.wavfile import write
|
||||
from env import AttrDict
|
||||
from meldataset import MAX_WAV_VALUE
|
||||
from bigvgan import BigVGAN as Generator
|
||||
|
||||
h = None
|
||||
device = None
|
||||
torch.backends.cudnn.benchmark = False
|
||||
|
||||
|
||||
def load_checkpoint(filepath, device):
|
||||
assert os.path.isfile(filepath)
|
||||
print(f"Loading '{filepath}'")
|
||||
checkpoint_dict = torch.load(filepath, map_location=device)
|
||||
print("Complete.")
|
||||
return checkpoint_dict
|
||||
|
||||
|
||||
def scan_checkpoint(cp_dir, prefix):
|
||||
pattern = os.path.join(cp_dir, prefix + "*")
|
||||
cp_list = glob.glob(pattern)
|
||||
if len(cp_list) == 0:
|
||||
return ""
|
||||
return sorted(cp_list)[-1]
|
||||
|
||||
|
||||
def inference(a, h):
|
||||
generator = Generator(h, use_cuda_kernel=a.use_cuda_kernel).to(device)
|
||||
|
||||
state_dict_g = load_checkpoint(a.checkpoint_file, device)
|
||||
generator.load_state_dict(state_dict_g["generator"])
|
||||
|
||||
filelist = os.listdir(a.input_mels_dir)
|
||||
|
||||
os.makedirs(a.output_dir, exist_ok=True)
|
||||
|
||||
generator.eval()
|
||||
generator.remove_weight_norm()
|
||||
with torch.no_grad():
|
||||
for i, filname in enumerate(filelist):
|
||||
# Load the mel spectrogram in .npy format
|
||||
x = np.load(os.path.join(a.input_mels_dir, filname))
|
||||
x = torch.FloatTensor(x).to(device)
|
||||
if len(x.shape) == 2:
|
||||
x = x.unsqueeze(0)
|
||||
|
||||
y_g_hat = generator(x)
|
||||
|
||||
audio = y_g_hat.squeeze()
|
||||
audio = audio * MAX_WAV_VALUE
|
||||
audio = audio.cpu().numpy().astype("int16")
|
||||
|
||||
output_file = os.path.join(a.output_dir, os.path.splitext(filname)[0] + "_generated_e2e.wav")
|
||||
write(output_file, h.sampling_rate, audio)
|
||||
print(output_file)
|
||||
|
||||
|
||||
def main():
|
||||
print("Initializing Inference Process..")
|
||||
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument("--input_mels_dir", default="test_mel_files")
|
||||
parser.add_argument("--output_dir", default="generated_files_from_mel")
|
||||
parser.add_argument("--checkpoint_file", required=True)
|
||||
parser.add_argument("--use_cuda_kernel", action="store_true", default=False)
|
||||
|
||||
a = parser.parse_args()
|
||||
|
||||
config_file = os.path.join(os.path.split(a.checkpoint_file)[0], "config.json")
|
||||
with open(config_file) as f:
|
||||
data = f.read()
|
||||
|
||||
global h
|
||||
json_config = json.loads(data)
|
||||
h = AttrDict(json_config)
|
||||
|
||||
torch.manual_seed(h.seed)
|
||||
global device
|
||||
if torch.cuda.is_available():
|
||||
torch.cuda.manual_seed(h.seed)
|
||||
device = torch.device("cuda")
|
||||
else:
|
||||
device = torch.device("cpu")
|
||||
|
||||
inference(a, h)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -1,238 +0,0 @@
|
||||
# Copyright (c) 2024 NVIDIA CORPORATION.
|
||||
# Licensed under the MIT license.
|
||||
|
||||
# Adapted from https://github.com/jik876/hifi-gan under the MIT license.
|
||||
# LICENSE is in incl_licenses directory.
|
||||
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
from librosa.filters import mel as librosa_mel_fn
|
||||
from scipy import signal
|
||||
|
||||
import typing
|
||||
from typing import List, Tuple
|
||||
from collections import namedtuple
|
||||
import math
|
||||
import functools
|
||||
|
||||
|
||||
# Adapted from https://github.com/descriptinc/descript-audio-codec/blob/main/dac/nn/loss.py under the MIT license.
|
||||
# LICENSE is in incl_licenses directory.
|
||||
class MultiScaleMelSpectrogramLoss(nn.Module):
|
||||
"""Compute distance between mel spectrograms. Can be used
|
||||
in a multi-scale way.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
n_mels : List[int]
|
||||
Number of mels per STFT, by default [5, 10, 20, 40, 80, 160, 320],
|
||||
window_lengths : List[int], optional
|
||||
Length of each window of each STFT, by default [32, 64, 128, 256, 512, 1024, 2048]
|
||||
loss_fn : typing.Callable, optional
|
||||
How to compare each loss, by default nn.L1Loss()
|
||||
clamp_eps : float, optional
|
||||
Clamp on the log magnitude, below, by default 1e-5
|
||||
mag_weight : float, optional
|
||||
Weight of raw magnitude portion of loss, by default 0.0 (no ampliciation on mag part)
|
||||
log_weight : float, optional
|
||||
Weight of log magnitude portion of loss, by default 1.0
|
||||
pow : float, optional
|
||||
Power to raise magnitude to before taking log, by default 1.0
|
||||
weight : float, optional
|
||||
Weight of this loss, by default 1.0
|
||||
match_stride : bool, optional
|
||||
Whether to match the stride of convolutional layers, by default False
|
||||
|
||||
Implementation copied from: https://github.com/descriptinc/lyrebird-audiotools/blob/961786aa1a9d628cca0c0486e5885a457fe70c1a/audiotools/metrics/spectral.py
|
||||
Additional code copied and modified from https://github.com/descriptinc/audiotools/blob/master/audiotools/core/audio_signal.py
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
sampling_rate: int,
|
||||
n_mels: List[int] = [5, 10, 20, 40, 80, 160, 320],
|
||||
window_lengths: List[int] = [32, 64, 128, 256, 512, 1024, 2048],
|
||||
loss_fn: typing.Callable = nn.L1Loss(),
|
||||
clamp_eps: float = 1e-5,
|
||||
mag_weight: float = 0.0,
|
||||
log_weight: float = 1.0,
|
||||
pow: float = 1.0,
|
||||
weight: float = 1.0,
|
||||
match_stride: bool = False,
|
||||
mel_fmin: List[float] = [0, 0, 0, 0, 0, 0, 0],
|
||||
mel_fmax: List[float] = [None, None, None, None, None, None, None],
|
||||
window_type: str = "hann",
|
||||
):
|
||||
super().__init__()
|
||||
self.sampling_rate = sampling_rate
|
||||
|
||||
STFTParams = namedtuple(
|
||||
"STFTParams",
|
||||
["window_length", "hop_length", "window_type", "match_stride"],
|
||||
)
|
||||
|
||||
self.stft_params = [
|
||||
STFTParams(
|
||||
window_length=w,
|
||||
hop_length=w // 4,
|
||||
match_stride=match_stride,
|
||||
window_type=window_type,
|
||||
)
|
||||
for w in window_lengths
|
||||
]
|
||||
self.n_mels = n_mels
|
||||
self.loss_fn = loss_fn
|
||||
self.clamp_eps = clamp_eps
|
||||
self.log_weight = log_weight
|
||||
self.mag_weight = mag_weight
|
||||
self.weight = weight
|
||||
self.mel_fmin = mel_fmin
|
||||
self.mel_fmax = mel_fmax
|
||||
self.pow = pow
|
||||
|
||||
@staticmethod
|
||||
@functools.lru_cache(None)
|
||||
def get_window(
|
||||
window_type,
|
||||
window_length,
|
||||
):
|
||||
return signal.get_window(window_type, window_length)
|
||||
|
||||
@staticmethod
|
||||
@functools.lru_cache(None)
|
||||
def get_mel_filters(sr, n_fft, n_mels, fmin, fmax):
|
||||
return librosa_mel_fn(sr=sr, n_fft=n_fft, n_mels=n_mels, fmin=fmin, fmax=fmax)
|
||||
|
||||
def mel_spectrogram(
|
||||
self,
|
||||
wav,
|
||||
n_mels,
|
||||
fmin,
|
||||
fmax,
|
||||
window_length,
|
||||
hop_length,
|
||||
match_stride,
|
||||
window_type,
|
||||
):
|
||||
"""
|
||||
Mirrors AudioSignal.mel_spectrogram used by BigVGAN-v2 training from:
|
||||
https://github.com/descriptinc/audiotools/blob/master/audiotools/core/audio_signal.py
|
||||
"""
|
||||
B, C, T = wav.shape
|
||||
|
||||
if match_stride:
|
||||
assert hop_length == window_length // 4, "For match_stride, hop must equal n_fft // 4"
|
||||
right_pad = math.ceil(T / hop_length) * hop_length - T
|
||||
pad = (window_length - hop_length) // 2
|
||||
else:
|
||||
right_pad = 0
|
||||
pad = 0
|
||||
|
||||
wav = torch.nn.functional.pad(wav, (pad, pad + right_pad), mode="reflect")
|
||||
|
||||
window = self.get_window(window_type, window_length)
|
||||
window = torch.from_numpy(window).to(wav.device).float()
|
||||
|
||||
stft = torch.stft(
|
||||
wav.reshape(-1, T),
|
||||
n_fft=window_length,
|
||||
hop_length=hop_length,
|
||||
window=window,
|
||||
return_complex=True,
|
||||
center=True,
|
||||
)
|
||||
_, nf, nt = stft.shape
|
||||
stft = stft.reshape(B, C, nf, nt)
|
||||
if match_stride:
|
||||
"""
|
||||
Drop first two and last two frames, which are added, because of padding. Now num_frames * hop_length = num_samples.
|
||||
"""
|
||||
stft = stft[..., 2:-2]
|
||||
magnitude = torch.abs(stft)
|
||||
|
||||
nf = magnitude.shape[2]
|
||||
mel_basis = self.get_mel_filters(self.sampling_rate, 2 * (nf - 1), n_mels, fmin, fmax)
|
||||
mel_basis = torch.from_numpy(mel_basis).to(wav.device)
|
||||
mel_spectrogram = magnitude.transpose(2, -1) @ mel_basis.T
|
||||
mel_spectrogram = mel_spectrogram.transpose(-1, 2)
|
||||
|
||||
return mel_spectrogram
|
||||
|
||||
def forward(self, x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
|
||||
"""Computes mel loss between an estimate and a reference
|
||||
signal.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
x : torch.Tensor
|
||||
Estimate signal
|
||||
y : torch.Tensor
|
||||
Reference signal
|
||||
|
||||
Returns
|
||||
-------
|
||||
torch.Tensor
|
||||
Mel loss.
|
||||
"""
|
||||
|
||||
loss = 0.0
|
||||
for n_mels, fmin, fmax, s in zip(self.n_mels, self.mel_fmin, self.mel_fmax, self.stft_params):
|
||||
kwargs = {
|
||||
"n_mels": n_mels,
|
||||
"fmin": fmin,
|
||||
"fmax": fmax,
|
||||
"window_length": s.window_length,
|
||||
"hop_length": s.hop_length,
|
||||
"match_stride": s.match_stride,
|
||||
"window_type": s.window_type,
|
||||
}
|
||||
|
||||
x_mels = self.mel_spectrogram(x, **kwargs)
|
||||
y_mels = self.mel_spectrogram(y, **kwargs)
|
||||
x_logmels = torch.log(x_mels.clamp(min=self.clamp_eps).pow(self.pow)) / torch.log(torch.tensor(10.0))
|
||||
y_logmels = torch.log(y_mels.clamp(min=self.clamp_eps).pow(self.pow)) / torch.log(torch.tensor(10.0))
|
||||
|
||||
loss += self.log_weight * self.loss_fn(x_logmels, y_logmels)
|
||||
loss += self.mag_weight * self.loss_fn(x_logmels, y_logmels)
|
||||
|
||||
return loss
|
||||
|
||||
|
||||
# Loss functions
|
||||
def feature_loss(fmap_r: List[List[torch.Tensor]], fmap_g: List[List[torch.Tensor]]) -> torch.Tensor:
|
||||
loss = 0
|
||||
for dr, dg in zip(fmap_r, fmap_g):
|
||||
for rl, gl in zip(dr, dg):
|
||||
loss += torch.mean(torch.abs(rl - gl))
|
||||
|
||||
return loss * 2 # This equates to lambda=2.0 for the feature matching loss
|
||||
|
||||
|
||||
def discriminator_loss(
|
||||
disc_real_outputs: List[torch.Tensor], disc_generated_outputs: List[torch.Tensor]
|
||||
) -> Tuple[torch.Tensor, List[torch.Tensor], List[torch.Tensor]]:
|
||||
loss = 0
|
||||
r_losses = []
|
||||
g_losses = []
|
||||
for dr, dg in zip(disc_real_outputs, disc_generated_outputs):
|
||||
r_loss = torch.mean((1 - dr) ** 2)
|
||||
g_loss = torch.mean(dg**2)
|
||||
loss += r_loss + g_loss
|
||||
r_losses.append(r_loss.item())
|
||||
g_losses.append(g_loss.item())
|
||||
|
||||
return loss, r_losses, g_losses
|
||||
|
||||
|
||||
def generator_loss(
|
||||
disc_outputs: List[torch.Tensor],
|
||||
) -> Tuple[torch.Tensor, List[torch.Tensor]]:
|
||||
loss = 0
|
||||
gen_losses = []
|
||||
for dg in disc_outputs:
|
||||
l = torch.mean((1 - dg) ** 2)
|
||||
gen_losses.append(l)
|
||||
loss += l
|
||||
|
||||
return loss, gen_losses
|
||||
@@ -1,370 +0,0 @@
|
||||
# Copyright (c) 2024 NVIDIA CORPORATION.
|
||||
# Licensed under the MIT license.
|
||||
|
||||
# Adapted from https://github.com/jik876/hifi-gan under the MIT license.
|
||||
# LICENSE is in incl_licenses directory.
|
||||
|
||||
import math
|
||||
import os
|
||||
import random
|
||||
import torch
|
||||
import torch.utils.data
|
||||
import numpy as np
|
||||
import librosa
|
||||
from librosa.filters import mel as librosa_mel_fn
|
||||
import pathlib
|
||||
from tqdm import tqdm
|
||||
from typing import List, Tuple, Optional
|
||||
from .env import AttrDict
|
||||
|
||||
MAX_WAV_VALUE = 32767.0 # NOTE: 32768.0 -1 to prevent int16 overflow (results in popping sound in corner cases)
|
||||
|
||||
|
||||
def dynamic_range_compression(x, C=1, clip_val=1e-5):
|
||||
return np.log(np.clip(x, a_min=clip_val, a_max=None) * C)
|
||||
|
||||
|
||||
def dynamic_range_decompression(x, C=1):
|
||||
return np.exp(x) / C
|
||||
|
||||
|
||||
def dynamic_range_compression_torch(x, C=1, clip_val=1e-5):
|
||||
return torch.log(torch.clamp(x, min=clip_val) * C)
|
||||
|
||||
|
||||
def dynamic_range_decompression_torch(x, C=1):
|
||||
return torch.exp(x) / C
|
||||
|
||||
|
||||
def spectral_normalize_torch(magnitudes):
|
||||
return dynamic_range_compression_torch(magnitudes)
|
||||
|
||||
|
||||
def spectral_de_normalize_torch(magnitudes):
|
||||
return dynamic_range_decompression_torch(magnitudes)
|
||||
|
||||
|
||||
mel_basis_cache = {}
|
||||
hann_window_cache = {}
|
||||
|
||||
|
||||
def mel_spectrogram(
|
||||
y: torch.Tensor,
|
||||
n_fft: int,
|
||||
num_mels: int,
|
||||
sampling_rate: int,
|
||||
hop_size: int,
|
||||
win_size: int,
|
||||
fmin: int,
|
||||
fmax: int = None,
|
||||
center: bool = False,
|
||||
) -> torch.Tensor:
|
||||
"""
|
||||
Calculate the mel spectrogram of an input signal.
|
||||
This function uses slaney norm for the librosa mel filterbank (using librosa.filters.mel) and uses Hann window for STFT (using torch.stft).
|
||||
|
||||
Args:
|
||||
y (torch.Tensor): Input signal.
|
||||
n_fft (int): FFT size.
|
||||
num_mels (int): Number of mel bins.
|
||||
sampling_rate (int): Sampling rate of the input signal.
|
||||
hop_size (int): Hop size for STFT.
|
||||
win_size (int): Window size for STFT.
|
||||
fmin (int): Minimum frequency for mel filterbank.
|
||||
fmax (int): Maximum frequency for mel filterbank. If None, defaults to half the sampling rate (fmax = sr / 2.0) inside librosa_mel_fn
|
||||
center (bool): Whether to pad the input to center the frames. Default is False.
|
||||
|
||||
Returns:
|
||||
torch.Tensor: Mel spectrogram.
|
||||
"""
|
||||
if torch.min(y) < -1.0:
|
||||
print(f"[WARNING] Min value of input waveform signal is {torch.min(y)}")
|
||||
if torch.max(y) > 1.0:
|
||||
print(f"[WARNING] Max value of input waveform signal is {torch.max(y)}")
|
||||
|
||||
device = y.device
|
||||
key = f"{n_fft}_{num_mels}_{sampling_rate}_{hop_size}_{win_size}_{fmin}_{fmax}_{device}"
|
||||
|
||||
if key not in mel_basis_cache:
|
||||
mel = librosa_mel_fn(sr=sampling_rate, n_fft=n_fft, n_mels=num_mels, fmin=fmin, fmax=fmax)
|
||||
mel_basis_cache[key] = torch.from_numpy(mel).float().to(device)
|
||||
hann_window_cache[key] = torch.hann_window(win_size).to(device)
|
||||
|
||||
mel_basis = mel_basis_cache[key]
|
||||
hann_window = hann_window_cache[key]
|
||||
|
||||
padding = (n_fft - hop_size) // 2
|
||||
y = torch.nn.functional.pad(y.unsqueeze(1), (padding, padding), mode="reflect").squeeze(1)
|
||||
|
||||
spec = torch.stft(
|
||||
y,
|
||||
n_fft,
|
||||
hop_length=hop_size,
|
||||
win_length=win_size,
|
||||
window=hann_window,
|
||||
center=center,
|
||||
pad_mode="reflect",
|
||||
normalized=False,
|
||||
onesided=True,
|
||||
return_complex=True,
|
||||
)
|
||||
spec = torch.sqrt(torch.view_as_real(spec).pow(2).sum(-1) + 1e-9)
|
||||
|
||||
mel_spec = torch.matmul(mel_basis, spec)
|
||||
mel_spec = spectral_normalize_torch(mel_spec)
|
||||
|
||||
return mel_spec
|
||||
|
||||
|
||||
def get_mel_spectrogram(wav, h):
|
||||
"""
|
||||
Generate mel spectrogram from a waveform using given hyperparameters.
|
||||
|
||||
Args:
|
||||
wav (torch.Tensor): Input waveform.
|
||||
h: Hyperparameters object with attributes n_fft, num_mels, sampling_rate, hop_size, win_size, fmin, fmax.
|
||||
|
||||
Returns:
|
||||
torch.Tensor: Mel spectrogram.
|
||||
"""
|
||||
return mel_spectrogram(
|
||||
wav,
|
||||
h.n_fft,
|
||||
h.num_mels,
|
||||
h.sampling_rate,
|
||||
h.hop_size,
|
||||
h.win_size,
|
||||
h.fmin,
|
||||
h.fmax,
|
||||
)
|
||||
|
||||
|
||||
def get_dataset_filelist(a):
|
||||
training_files = []
|
||||
validation_files = []
|
||||
list_unseen_validation_files = []
|
||||
|
||||
with open(a.input_training_file, "r", encoding="utf-8") as fi:
|
||||
training_files = [
|
||||
os.path.join(a.input_wavs_dir, x.split("|")[0] + ".wav") for x in fi.read().split("\n") if len(x) > 0
|
||||
]
|
||||
print(f"first training file: {training_files[0]}")
|
||||
|
||||
with open(a.input_validation_file, "r", encoding="utf-8") as fi:
|
||||
validation_files = [
|
||||
os.path.join(a.input_wavs_dir, x.split("|")[0] + ".wav") for x in fi.read().split("\n") if len(x) > 0
|
||||
]
|
||||
print(f"first validation file: {validation_files[0]}")
|
||||
|
||||
for i in range(len(a.list_input_unseen_validation_file)):
|
||||
with open(a.list_input_unseen_validation_file[i], "r", encoding="utf-8") as fi:
|
||||
unseen_validation_files = [
|
||||
os.path.join(a.list_input_unseen_wavs_dir[i], x.split("|")[0] + ".wav")
|
||||
for x in fi.read().split("\n")
|
||||
if len(x) > 0
|
||||
]
|
||||
print(f"first unseen {i}th validation fileset: {unseen_validation_files[0]}")
|
||||
list_unseen_validation_files.append(unseen_validation_files)
|
||||
|
||||
return training_files, validation_files, list_unseen_validation_files
|
||||
|
||||
|
||||
class MelDataset(torch.utils.data.Dataset):
|
||||
def __init__(
|
||||
self,
|
||||
training_files: List[str],
|
||||
hparams: AttrDict,
|
||||
segment_size: int,
|
||||
n_fft: int,
|
||||
num_mels: int,
|
||||
hop_size: int,
|
||||
win_size: int,
|
||||
sampling_rate: int,
|
||||
fmin: int,
|
||||
fmax: Optional[int],
|
||||
split: bool = True,
|
||||
shuffle: bool = True,
|
||||
device: str = None,
|
||||
fmax_loss: Optional[int] = None,
|
||||
fine_tuning: bool = False,
|
||||
base_mels_path: str = None,
|
||||
is_seen: bool = True,
|
||||
):
|
||||
self.audio_files = training_files
|
||||
random.seed(1234)
|
||||
if shuffle:
|
||||
random.shuffle(self.audio_files)
|
||||
self.hparams = hparams
|
||||
self.is_seen = is_seen
|
||||
if self.is_seen:
|
||||
self.name = pathlib.Path(self.audio_files[0]).parts[0]
|
||||
else:
|
||||
self.name = "-".join(pathlib.Path(self.audio_files[0]).parts[:2]).strip("/")
|
||||
|
||||
self.segment_size = segment_size
|
||||
self.sampling_rate = sampling_rate
|
||||
self.split = split
|
||||
self.n_fft = n_fft
|
||||
self.num_mels = num_mels
|
||||
self.hop_size = hop_size
|
||||
self.win_size = win_size
|
||||
self.fmin = fmin
|
||||
self.fmax = fmax
|
||||
self.fmax_loss = fmax_loss
|
||||
self.device = device
|
||||
self.fine_tuning = fine_tuning
|
||||
self.base_mels_path = base_mels_path
|
||||
|
||||
print("[INFO] checking dataset integrity...")
|
||||
for i in tqdm(range(len(self.audio_files))):
|
||||
assert os.path.exists(self.audio_files[i]), f"{self.audio_files[i]} not found"
|
||||
|
||||
def __getitem__(self, index: int) -> Tuple[torch.Tensor, torch.Tensor, str, torch.Tensor]:
|
||||
try:
|
||||
filename = self.audio_files[index]
|
||||
|
||||
# Use librosa.load that ensures loading waveform into mono with [-1, 1] float values
|
||||
# Audio is ndarray with shape [T_time]. Disable auto-resampling here to minimize overhead
|
||||
# The on-the-fly resampling during training will be done only for the obtained random chunk
|
||||
audio, source_sampling_rate = librosa.load(filename, sr=None, mono=True)
|
||||
|
||||
# Main logic that uses <mel, audio> pair for training BigVGAN
|
||||
if not self.fine_tuning:
|
||||
if self.split: # Training step
|
||||
# Obtain randomized audio chunk
|
||||
if source_sampling_rate != self.sampling_rate:
|
||||
# Adjust segment size to crop if the source sr is different
|
||||
target_segment_size = math.ceil(self.segment_size * (source_sampling_rate / self.sampling_rate))
|
||||
else:
|
||||
target_segment_size = self.segment_size
|
||||
|
||||
# Compute upper bound index for the random chunk
|
||||
random_chunk_upper_bound = max(0, audio.shape[0] - target_segment_size)
|
||||
|
||||
# Crop or pad audio to obtain random chunk with target_segment_size
|
||||
if audio.shape[0] >= target_segment_size:
|
||||
audio_start = random.randint(0, random_chunk_upper_bound)
|
||||
audio = audio[audio_start : audio_start + target_segment_size]
|
||||
else:
|
||||
audio = np.pad(
|
||||
audio,
|
||||
(0, target_segment_size - audio.shape[0]),
|
||||
mode="constant",
|
||||
)
|
||||
|
||||
# Resample audio chunk to self.sampling rate
|
||||
if source_sampling_rate != self.sampling_rate:
|
||||
audio = librosa.resample(
|
||||
audio,
|
||||
orig_sr=source_sampling_rate,
|
||||
target_sr=self.sampling_rate,
|
||||
)
|
||||
if audio.shape[0] > self.segment_size:
|
||||
# trim last elements to match self.segment_size (e.g., 16385 for 44khz downsampled to 24khz -> 16384)
|
||||
audio = audio[: self.segment_size]
|
||||
|
||||
else: # Validation step
|
||||
# Resample full audio clip to target sampling rate
|
||||
if source_sampling_rate != self.sampling_rate:
|
||||
audio = librosa.resample(
|
||||
audio,
|
||||
orig_sr=source_sampling_rate,
|
||||
target_sr=self.sampling_rate,
|
||||
)
|
||||
# Trim last elements to match audio length to self.hop_size * n for evaluation
|
||||
if (audio.shape[0] % self.hop_size) != 0:
|
||||
audio = audio[: -(audio.shape[0] % self.hop_size)]
|
||||
|
||||
# BigVGAN is trained using volume-normalized waveform
|
||||
audio = librosa.util.normalize(audio) * 0.95
|
||||
|
||||
# Cast ndarray to torch tensor
|
||||
audio = torch.FloatTensor(audio)
|
||||
audio = audio.unsqueeze(0) # [B(1), self.segment_size]
|
||||
|
||||
# Compute mel spectrogram corresponding to audio
|
||||
mel = mel_spectrogram(
|
||||
audio,
|
||||
self.n_fft,
|
||||
self.num_mels,
|
||||
self.sampling_rate,
|
||||
self.hop_size,
|
||||
self.win_size,
|
||||
self.fmin,
|
||||
self.fmax,
|
||||
center=False,
|
||||
) # [B(1), self.num_mels, self.segment_size // self.hop_size]
|
||||
|
||||
# Fine-tuning logic that uses pre-computed mel. Example: Using TTS model-generated mel as input
|
||||
else:
|
||||
# For fine-tuning, assert that the waveform is in the defined sampling_rate
|
||||
# Fine-tuning won't support on-the-fly resampling to be fool-proof (the dataset should have been prepared properly)
|
||||
assert source_sampling_rate == self.sampling_rate, (
|
||||
f"For fine_tuning, waveform must be in the spcified sampling rate {self.sampling_rate}, got {source_sampling_rate}"
|
||||
)
|
||||
|
||||
# Cast ndarray to torch tensor
|
||||
audio = torch.FloatTensor(audio)
|
||||
audio = audio.unsqueeze(0) # [B(1), T_time]
|
||||
|
||||
# Load pre-computed mel from disk
|
||||
mel = np.load(
|
||||
os.path.join(
|
||||
self.base_mels_path,
|
||||
os.path.splitext(os.path.split(filename)[-1])[0] + ".npy",
|
||||
)
|
||||
)
|
||||
mel = torch.from_numpy(mel)
|
||||
|
||||
if len(mel.shape) < 3:
|
||||
mel = mel.unsqueeze(0) # ensure [B, C, T]
|
||||
|
||||
if self.split:
|
||||
frames_per_seg = math.ceil(self.segment_size / self.hop_size)
|
||||
|
||||
if audio.size(1) >= self.segment_size:
|
||||
mel_start = random.randint(0, mel.size(2) - frames_per_seg - 1)
|
||||
mel = mel[:, :, mel_start : mel_start + frames_per_seg]
|
||||
audio = audio[
|
||||
:,
|
||||
mel_start * self.hop_size : (mel_start + frames_per_seg) * self.hop_size,
|
||||
]
|
||||
|
||||
# Pad pre-computed mel and audio to match length to ensuring fine-tuning without error.
|
||||
# NOTE: this may introduce a single-frame misalignment of the <pre-computed mel, audio>
|
||||
# To remove possible misalignment, it is recommended to prepare the <pre-computed mel, audio> pair where the audio length is the integer multiple of self.hop_size
|
||||
mel = torch.nn.functional.pad(mel, (0, frames_per_seg - mel.size(2)), "constant")
|
||||
audio = torch.nn.functional.pad(audio, (0, self.segment_size - audio.size(1)), "constant")
|
||||
|
||||
# Compute mel_loss used by spectral regression objective. Uses self.fmax_loss instead (usually None)
|
||||
mel_loss = mel_spectrogram(
|
||||
audio,
|
||||
self.n_fft,
|
||||
self.num_mels,
|
||||
self.sampling_rate,
|
||||
self.hop_size,
|
||||
self.win_size,
|
||||
self.fmin,
|
||||
self.fmax_loss,
|
||||
center=False,
|
||||
) # [B(1), self.num_mels, self.segment_size // self.hop_size]
|
||||
|
||||
# Shape sanity checks
|
||||
assert (
|
||||
audio.shape[1] == mel.shape[2] * self.hop_size and audio.shape[1] == mel_loss.shape[2] * self.hop_size
|
||||
), (
|
||||
f"Audio length must be mel frame length * hop_size. Got audio shape {audio.shape} mel shape {mel.shape} mel_loss shape {mel_loss.shape}"
|
||||
)
|
||||
|
||||
return (mel.squeeze(), audio.squeeze(0), filename, mel_loss.squeeze())
|
||||
|
||||
# If it encounters error during loading the data, skip this sample and load random other sample to the batch
|
||||
except Exception as e:
|
||||
if self.fine_tuning:
|
||||
raise e # Terminate training if it is fine-tuning. The dataset should have been prepared properly.
|
||||
else:
|
||||
print(f"[WARNING] Failed to load waveform, skipping! filename: {filename} Error: {e}")
|
||||
return self[random.randrange(len(self))]
|
||||
|
||||
def __len__(self):
|
||||
return len(self.audio_files)
|
||||
@@ -1 +0,0 @@
|
||||
|
||||
@@ -1,4 +0,0 @@
|
||||
| Field | Response |
|
||||
| :--------------------------------------------------------------------------------------------------------- | :--------------------------------------------------- |
|
||||
| Participation considerations from adversely impacted groups protected classes in model design and testing: | None |
|
||||
| Measures taken to mitigate against unwanted bias: | No measures taken to mitigate against unwanted bias. |
|
||||
@@ -1,13 +0,0 @@
|
||||
| Field | Response |
|
||||
| :---------------------------------------------------------------------------------------------------- | :--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- |
|
||||
| Intended Application & Domain: | Generating waveform from mel spectrogram. |
|
||||
| Model Type: | Convolutional Neural Network (CNN) |
|
||||
| Intended Users: | This model is intended for developers to synthesize and generate waveforms from the AI-generated mel spectrograms. |
|
||||
| Output: | Audio Waveform |
|
||||
| Describe how the model works: | Model generates audio waveform corresponding to the input mel spectrogram. |
|
||||
| Name the adversely impacted groups this has been tested to deliver comparable outcomes regardless of: | Not Applicable |
|
||||
| Technical Limitations: | This may not perform well on synthetically-generated mel spectrograms that deviate significantly from the profile of mel spectrograms on which this was trained. |
|
||||
| Verified to have met prescribed NVIDIA quality standards: | Yes |
|
||||
| Performance Metrics: | Perceptual Evaluation of Speech Quality (PESQ), Virtual Speech Quality Objective Listener (VISQOL), Multi-resolution STFT (MRSTFT), Mel cepstral distortion (MCD), Periodicity RMSE, Voice/Unvoiced F1 Score (V/UV F1) |
|
||||
| Potential Known Risks: | This model may generate low-quality or distorted soundwaves. |
|
||||
| Licensing: | https://github.com/NVIDIA/BigVGAN/blob/main/LICENSE |
|
||||
@@ -1,126 +0,0 @@
|
||||
# Model Overview
|
||||
|
||||
## Description:
|
||||
|
||||
BigVGAN is a generative AI model specialized in synthesizing audio waveforms using Mel spectrogram as inputs.
|
||||
|
||||
<center><img src="https://user-images.githubusercontent.com/15963413/218609148-881e39df-33af-4af9-ab95-1427c4ebf062.png" width="800"></center>
|
||||
|
||||
BigVGAN is a fully convolutional architecture with several upsampling blocks using transposed convolution followed by multiple residual dilated convolution layers.
|
||||
|
||||
BigVGAN consists of a novel module, called anti-aliased multi-periodicity composition (AMP), which is specifically designed for generating waveforms. AMP is specialized in synthesizing high-frequency and periodic soundwaves drawing inspiration from audio signal processing principles.
|
||||
|
||||
It applies a periodic activation function, called Snake, which provides an inductive bias to the architecture in generating periodic soundwaves. It also applies anti-aliasing filters to reduce undesired artifacts in the generated waveforms. <br>
|
||||
|
||||
This model is ready for commercial use.<br>
|
||||
|
||||
## References(s):
|
||||
|
||||
- [BigVGAN: A Universal Neural Vocoder with Large-Scale Training](https://arxiv.org/abs/2206.04658) <br>
|
||||
- [Project Page](https://research.nvidia.com/labs/adlr/projects/bigvgan/) <br>
|
||||
- [Audio Demo](https://bigvgan-demo.github.io/) <br>
|
||||
|
||||
## Model Architecture:
|
||||
|
||||
**Architecture Type:** Convolution Neural Network (CNN) <br>
|
||||
**Network Architecture:** You can see the details of this model on this link: https://github.com/NVIDIA/BigVGAN and the related paper can be found here: https://arxiv.org/abs/2206.04658<br>
|
||||
**Model Version:** 2.0 <br>
|
||||
|
||||
## Input:
|
||||
|
||||
**Input Type:** Audio <br>
|
||||
**Input Format:** Mel Spectrogram <br>
|
||||
**Input Parameters:** None <br>
|
||||
**Other Properties Related to Input:** The input mel spectrogram has shape `[batch, channels, frames]`, where `channels` refers to the number of mel bands defined by the model and `frames` refers to the temporal length. The model supports arbitrary long `frames` that fits into the GPU memory.
|
||||
|
||||
## Output:
|
||||
|
||||
**Input Type:** Audio <br>
|
||||
**Output Format:** Audio Waveform <br>
|
||||
**Output Parameters:** None <br>
|
||||
**Other Properties Related to Output:** The output audio waveform has shape `[batch, 1, time]`, where `1` refers to the mono audio channels and `time` refers to the temporal length. `time` is defined as a fixed integer multiple of input `frames`, which is an upsampling ratio of the model (`time = upsampling ratio * frames`). The output audio waveform consitutes float values with a range of `[-1, 1]`.
|
||||
|
||||
## Software Integration:
|
||||
|
||||
**Runtime Engine(s):** PyTorch
|
||||
|
||||
**Supported Hardware Microarchitecture Compatibility:** NVIDIA Ampere, NVIDIA Hopper, NVIDIA Lovelace, NVIDIA Turing, NVIDIA Volta <br>
|
||||
|
||||
## Preferred/Supported Operating System(s):
|
||||
|
||||
Linux
|
||||
|
||||
## Model Version(s):
|
||||
|
||||
v2.0
|
||||
|
||||
## Training, Testing, and Evaluation Datasets:
|
||||
|
||||
### Training Dataset:
|
||||
|
||||
The dataset contains diverse audio types, including speech in multiple languages, environmental sounds, and instruments.
|
||||
|
||||
**Links:**
|
||||
|
||||
- [AAM: Artificial Audio Multitracks Dataset](https://zenodo.org/records/5794629)
|
||||
- [AudioCaps](https://audiocaps.github.io/)
|
||||
- [AudioSet](https://research.google.com/audioset/index.html)
|
||||
- [common-accent](https://huggingface.co/datasets/DTU54DL/common-accent)
|
||||
- [Crowd Sourced Emotional Multimodal Actors Dataset (CREMA-D)](https://ieeexplore.ieee.org/document/6849440)
|
||||
- [DCASE2017 Challenge, Task 4: Large-scale weakly supervised sound event detection for smart cars](https://dcase.community/challenge2017/task-large-scale-sound-event-detection)
|
||||
- [FSDnoisy18k](https://zenodo.org/records/2529934)
|
||||
- [Free Universal Sound Separation Dataset](https://zenodo.org/records/3694384)
|
||||
- [Greatest Hits dataset](https://andrewowens.com/vis/)
|
||||
- [GTZAN](https://ieeexplore.ieee.org/document/1021072)
|
||||
- [JL corpus](https://www.kaggle.com/datasets/tli725/jl-corpus)
|
||||
- [Medley-solos-DB: a cross-collection dataset for musical instrument recognition](https://zenodo.org/records/3464194)
|
||||
- [MUSAN: A Music, Speech, and Noise Corpus](https://www.openslr.org/17/)
|
||||
- [MusicBench](https://huggingface.co/datasets/amaai-lab/MusicBench)
|
||||
- [MusicCaps](https://www.kaggle.com/datasets/googleai/musiccaps)
|
||||
- [MusicNet](https://www.kaggle.com/datasets/imsparsh/musicnet-dataset)
|
||||
- [NSynth](https://magenta.tensorflow.org/datasets/nsynth)
|
||||
- [OnAir-Music-Dataset](https://github.com/sevagh/OnAir-Music-Dataset)
|
||||
- [Audio Piano Triads Dataset](https://zenodo.org/records/4740877)
|
||||
- [Pitch Audio Dataset (Surge synthesizer)](https://zenodo.org/records/4677097)
|
||||
- [SONYC Urban Sound Tagging (SONYC-UST): a multilabel dataset from an urban acoustic sensor network](https://zenodo.org/records/3966543)
|
||||
- [VocalSound: A Dataset for Improving Human Vocal Sounds Recognition](https://arxiv.org/abs/2205.03433)
|
||||
- [WavText5K](https://github.com/microsoft/WavText5K)
|
||||
- [CSS10: A Collection of Single Speaker Speech Datasets for 10 Languages](https://github.com/Kyubyong/css10)
|
||||
- [Hi-Fi Multi-Speaker English TTS Dataset (Hi-Fi TTS)](https://www.openslr.org/109/)
|
||||
- [IIIT-H Indic Speech Databases](http://festvox.org/databases/iiit_voices/)
|
||||
- [Libri-Light: A Benchmark for ASR with Limited or No Supervision](https://arxiv.org/abs/1912.07875)
|
||||
- [LibriTTS: A Corpus Derived from LibriSpeech for Text-to-Speech](https://www.openslr.org/60)
|
||||
- [LibriTTS-R: A Restored Multi-Speaker Text-to-Speech Corpus](https://www.openslr.org/141/)
|
||||
- [The SIWIS French Speech Synthesis Database](https://datashare.ed.ac.uk/handle/10283/2353)
|
||||
- [Crowdsourced high-quality Colombian Spanish speech data set](https://openslr.org/72/)
|
||||
- [TTS-Portuguese Corpus](https://github.com/Edresson/TTS-Portuguese-Corpus)
|
||||
- [CSTR VCTK Corpus: English Multi-speaker Corpus for CSTR Voice Cloning Toolkit](https://datashare.ed.ac.uk/handle/10283/3443)
|
||||
|
||||
\*\* Data Collection Method by dataset <br>
|
||||
|
||||
- Human <br>
|
||||
|
||||
\*\* Labeling Method by dataset (for those with labels) <br>
|
||||
|
||||
- Hybrid: Automated, Human, Unknown <br>
|
||||
|
||||
### Evaluating Dataset:
|
||||
|
||||
Properties: The audio generation quality of BigVGAN is evaluated using `dev` splits of the [LibriTTS dataset](https://www.openslr.org/60/) and [Hi-Fi TTS dataset](https://www.openslr.org/109/). The datasets include speech in English language with equal balance of genders.
|
||||
|
||||
\*\* Data Collection Method by dataset <br>
|
||||
|
||||
- Human <br>
|
||||
|
||||
\*\* Labeling Method by dataset <br>
|
||||
|
||||
- Automated <br>
|
||||
|
||||
## Inference:
|
||||
|
||||
**Engine:** PyTorch <br>
|
||||
**Test Hardware:** NVIDIA A100 GPU <br>
|
||||
|
||||
## Ethical Considerations:
|
||||
|
||||
NVIDIA believes Trustworthy AI is a shared responsibility and we have established policies and practices to enable development for a wide array of AI applications. When downloaded or used in accordance with our terms of service, developers should work with their internal model team to ensure this model meets requirements for the relevant industry and use case and addresses unforeseen product misuse. For more detailed information on ethical considerations for this model, please see the Model Card++ Explainability, Bias, Safety & Security, and Privacy Subcards. Please report security vulnerabilities or NVIDIA AI Concerns [here](https://www.nvidia.com/en-us/support/submit-security-vulnerability/).
|
||||
@@ -1,14 +0,0 @@
|
||||
| Field | Response |
|
||||
| :------------------------------------------------------------------------------------------------------------------------------------- | :--------------------------------------------- |
|
||||
| Generatable or reverse engineerable personal information? | None |
|
||||
| Protected class data used to create this model? | None |
|
||||
| Was consent obtained for any personal data used? | Not Applicable (No Personal Data) |
|
||||
| How often is dataset reviewed? | Before Release |
|
||||
| Is a mechanism in place to honor data subject right of access or deletion of personal data? | Not Applicable |
|
||||
| If personal collected for the development of the model, was it collected directly by NVIDIA? | Not Applicable |
|
||||
| If personal collected for the development of the model by NVIDIA, do you maintain or have access to disclosures made to data subjects? | Not Applicable |
|
||||
| If personal collected for the development of this AI model, was it minimized to only what was required? | Not Applicable |
|
||||
| Is data in dataset traceable? | Yes |
|
||||
| Is there provenance for all datasets used in training? | Yes |
|
||||
| Does data labeling (annotation, metadata) comply with privacy laws? | Yes |
|
||||
| Is data compliant with data subject requests for data correction or removal, if such a request was made? | No, not possible with externally-sourced data. |
|
||||
@@ -1,6 +0,0 @@
|
||||
| Field | Response |
|
||||
| :---------------------------------------------- | :---------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- |
|
||||
| Model Application(s): | Synethic Audio Generation |
|
||||
| Describe the life critical impact (if present). | Not Applicable |
|
||||
| Use Case Restrictions: | None |
|
||||
| Model and dataset restrictions: | The Principle of least privilege (PoLP) is applied limiting access for dataset generation and model development. Restrictions enforce dataset access during training, and dataset license constraints adhered to. |
|
||||
@@ -1,13 +0,0 @@
|
||||
torch
|
||||
numpy
|
||||
librosa>=0.8.1
|
||||
scipy
|
||||
tensorboard
|
||||
soundfile
|
||||
matplotlib
|
||||
pesq
|
||||
auraloss
|
||||
tqdm
|
||||
nnAudio
|
||||
ninja
|
||||
huggingface_hub>=0.23.4
|
||||
@@ -1,62 +0,0 @@
|
||||
# Copyright (c) 2024 NVIDIA CORPORATION.
|
||||
# Licensed under the MIT license.
|
||||
|
||||
import os
|
||||
import sys
|
||||
|
||||
# to import modules from parent_dir
|
||||
parent_dir = os.path.abspath(os.path.join(os.path.dirname(__file__), ".."))
|
||||
sys.path.append(parent_dir)
|
||||
|
||||
import torch
|
||||
from alias_free_activation.cuda import activation1d
|
||||
from activations import Snake
|
||||
|
||||
|
||||
def test_load_fused_kernels():
|
||||
try:
|
||||
print("[Success] load_fused_kernels")
|
||||
except ImportError as e:
|
||||
print("[Fail] load_fused_kernels")
|
||||
raise e
|
||||
|
||||
|
||||
def test_anti_alias_activation():
|
||||
data = torch.rand((10, 10, 200), device="cuda")
|
||||
|
||||
# Check activations.Snake cuda vs. torch
|
||||
fused_anti_alias_activation = activation1d.Activation1d(activation=Snake(10), fused=True).cuda()
|
||||
fused_activation_output = fused_anti_alias_activation(data)
|
||||
|
||||
torch_anti_alias_activation = activation1d.Activation1d(activation=Snake(10), fused=False).cuda()
|
||||
torch_activation_output = torch_anti_alias_activation(data)
|
||||
|
||||
test_result = (fused_activation_output - torch_activation_output).abs()
|
||||
|
||||
while test_result.dim() != 1:
|
||||
test_result = test_result.mean(dim=-1)
|
||||
|
||||
diff = test_result.mean(dim=-1)
|
||||
|
||||
if diff <= 1e-3:
|
||||
print(
|
||||
f"\n[Success] test_fused_anti_alias_activation"
|
||||
f"\n > mean_difference={diff}"
|
||||
f"\n > fused_values={fused_activation_output[-1][-1][:].tolist()}"
|
||||
f"\n > torch_values={torch_activation_output[-1][-1][:].tolist()}"
|
||||
)
|
||||
else:
|
||||
print(
|
||||
f"\n[Fail] test_fused_anti_alias_activation"
|
||||
f"\n > mean_difference={diff}, "
|
||||
f"\n > fused_values={fused_activation_output[-1][-1][:].tolist()}, "
|
||||
f"\n > torch_values={torch_activation_output[-1][-1][:].tolist()}"
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
from alias_free_activation.cuda import load
|
||||
|
||||
load.load()
|
||||
test_load_fused_kernels()
|
||||
test_anti_alias_activation()
|
||||
@@ -1,62 +0,0 @@
|
||||
# Copyright (c) 2024 NVIDIA CORPORATION.
|
||||
# Licensed under the MIT license.
|
||||
|
||||
import os
|
||||
import sys
|
||||
|
||||
# to import modules from parent_dir
|
||||
parent_dir = os.path.abspath(os.path.join(os.path.dirname(__file__), ".."))
|
||||
sys.path.append(parent_dir)
|
||||
|
||||
import torch
|
||||
from alias_free_activation.cuda import activation1d
|
||||
from activations import SnakeBeta
|
||||
|
||||
|
||||
def test_load_fused_kernels():
|
||||
try:
|
||||
print("[Success] load_fused_kernels")
|
||||
except ImportError as e:
|
||||
print("[Fail] load_fused_kernels")
|
||||
raise e
|
||||
|
||||
|
||||
def test_anti_alias_activation():
|
||||
data = torch.rand((10, 10, 200), device="cuda")
|
||||
|
||||
# Check activations, Snake CUDA vs. Torch
|
||||
fused_anti_alias_activation = activation1d.Activation1d(activation=SnakeBeta(10), fused=True).cuda()
|
||||
fused_activation_output = fused_anti_alias_activation(data)
|
||||
|
||||
torch_anti_alias_activation = activation1d.Activation1d(activation=SnakeBeta(10), fused=False).cuda()
|
||||
torch_activation_output = torch_anti_alias_activation(data)
|
||||
|
||||
test_result = (fused_activation_output - torch_activation_output).abs()
|
||||
|
||||
while test_result.dim() != 1:
|
||||
test_result = test_result.mean(dim=-1)
|
||||
|
||||
diff = test_result.mean(dim=-1)
|
||||
|
||||
if diff <= 1e-3:
|
||||
print(
|
||||
f"\n[Success] test_fused_anti_alias_activation"
|
||||
f"\n > mean_difference={diff}"
|
||||
f"\n > fused_values={fused_activation_output[-1][-1][:].tolist()}"
|
||||
f"\n > torch_values={torch_activation_output[-1][-1][:].tolist()}"
|
||||
)
|
||||
else:
|
||||
print(
|
||||
f"\n[Fail] test_fused_anti_alias_activation"
|
||||
f"\n > mean_difference={diff}, "
|
||||
f"\n > fused_values={fused_activation_output[-1][-1][:].tolist()}, "
|
||||
f"\n > torch_values={torch_activation_output[-1][-1][:].tolist()}"
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
from alias_free_activation.cuda import load
|
||||
|
||||
load.load()
|
||||
test_load_fused_kernels()
|
||||
test_anti_alias_activation()
|
||||
@@ -1,215 +0,0 @@
|
||||
# Copyright (c) 2024 NVIDIA CORPORATION.
|
||||
# Licensed under the MIT license.
|
||||
|
||||
import os
|
||||
import sys
|
||||
|
||||
# to import modules from parent_dir
|
||||
parent_dir = os.path.abspath(os.path.join(os.path.dirname(__file__), ".."))
|
||||
sys.path.append(parent_dir)
|
||||
|
||||
import torch
|
||||
import json
|
||||
from env import AttrDict
|
||||
from bigvgan import BigVGAN
|
||||
from time import time
|
||||
from tqdm import tqdm
|
||||
from meldataset import mel_spectrogram, MAX_WAV_VALUE
|
||||
from scipy.io.wavfile import write
|
||||
import numpy as np
|
||||
|
||||
import argparse
|
||||
|
||||
torch.backends.cudnn.benchmark = True
|
||||
|
||||
# For easier debugging
|
||||
torch.set_printoptions(linewidth=200, threshold=10_000)
|
||||
|
||||
|
||||
def generate_soundwave(duration=5.0, sr=24000):
|
||||
t = np.linspace(0, duration, int(sr * duration), False, dtype=np.float32)
|
||||
|
||||
modulation = np.sin(2 * np.pi * t / duration)
|
||||
|
||||
min_freq = 220
|
||||
max_freq = 1760
|
||||
frequencies = min_freq + (max_freq - min_freq) * (modulation + 1) / 2
|
||||
soundwave = np.sin(2 * np.pi * frequencies * t)
|
||||
|
||||
soundwave = soundwave / np.max(np.abs(soundwave)) * 0.95
|
||||
|
||||
return soundwave, sr
|
||||
|
||||
|
||||
def get_mel(x, h):
|
||||
return mel_spectrogram(x, h.n_fft, h.num_mels, h.sampling_rate, h.hop_size, h.win_size, h.fmin, h.fmax)
|
||||
|
||||
|
||||
def load_checkpoint(filepath, device):
|
||||
assert os.path.isfile(filepath)
|
||||
print(f"Loading '{filepath}'")
|
||||
checkpoint_dict = torch.load(filepath, map_location=device)
|
||||
print("Complete.")
|
||||
return checkpoint_dict
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
parser = argparse.ArgumentParser(description="Test script to check CUDA kernel correctness.")
|
||||
parser.add_argument(
|
||||
"--checkpoint_file",
|
||||
type=str,
|
||||
required=True,
|
||||
help="Path to the checkpoint file. Assumes config.json exists in the directory.",
|
||||
)
|
||||
|
||||
args = parser.parse_args()
|
||||
|
||||
config_file = os.path.join(os.path.split(args.checkpoint_file)[0], "config.json")
|
||||
with open(config_file) as f:
|
||||
config = f.read()
|
||||
json_config = json.loads(config)
|
||||
h = AttrDict({**json_config})
|
||||
|
||||
print("loading plain Pytorch BigVGAN")
|
||||
generator_original = BigVGAN(h).to("cuda")
|
||||
print("loading CUDA kernel BigVGAN with auto-build")
|
||||
generator_cuda_kernel = BigVGAN(h, use_cuda_kernel=True).to("cuda")
|
||||
|
||||
state_dict_g = load_checkpoint(args.checkpoint_file, "cuda")
|
||||
generator_original.load_state_dict(state_dict_g["generator"])
|
||||
generator_cuda_kernel.load_state_dict(state_dict_g["generator"])
|
||||
|
||||
generator_original.remove_weight_norm()
|
||||
generator_original.eval()
|
||||
generator_cuda_kernel.remove_weight_norm()
|
||||
generator_cuda_kernel.eval()
|
||||
|
||||
# define number of samples and length of mel frame to benchmark
|
||||
num_sample = 10
|
||||
num_mel_frame = 16384
|
||||
|
||||
# CUDA kernel correctness check
|
||||
diff = 0.0
|
||||
for i in tqdm(range(num_sample)):
|
||||
# Random mel
|
||||
data = torch.rand((1, h.num_mels, num_mel_frame), device="cuda")
|
||||
|
||||
with torch.inference_mode():
|
||||
audio_original = generator_original(data)
|
||||
|
||||
with torch.inference_mode():
|
||||
audio_cuda_kernel = generator_cuda_kernel(data)
|
||||
|
||||
# Both outputs should be (almost) the same
|
||||
test_result = (audio_original - audio_cuda_kernel).abs()
|
||||
diff += test_result.mean(dim=-1).item()
|
||||
|
||||
diff /= num_sample
|
||||
if diff <= 2e-3: # We can expect a small difference (~1e-3) which does not affect perceptual quality
|
||||
print(
|
||||
f"\n[Success] test CUDA fused vs. plain torch BigVGAN inference"
|
||||
f"\n > mean_difference={diff}"
|
||||
f"\n > fused_values={audio_cuda_kernel[-1][-1][-30:].tolist()}"
|
||||
f"\n > torch_values={audio_original[-1][-1][-30:].tolist()}"
|
||||
)
|
||||
else:
|
||||
print(
|
||||
f"\n[Fail] test CUDA fused vs. plain torch BigVGAN inference"
|
||||
f"\n > mean_difference={diff}"
|
||||
f"\n > fused_values={audio_cuda_kernel[-1][-1][-30:].tolist()}, "
|
||||
f"\n > torch_values={audio_original[-1][-1][-30:].tolist()}"
|
||||
)
|
||||
|
||||
del data, audio_original, audio_cuda_kernel
|
||||
|
||||
# Variables for tracking total time and VRAM usage
|
||||
toc_total_original = 0
|
||||
toc_total_cuda_kernel = 0
|
||||
vram_used_original_total = 0
|
||||
vram_used_cuda_kernel_total = 0
|
||||
audio_length_total = 0
|
||||
|
||||
# Measure Original inference in isolation
|
||||
for i in tqdm(range(num_sample)):
|
||||
torch.cuda.reset_peak_memory_stats(device="cuda")
|
||||
data = torch.rand((1, h.num_mels, num_mel_frame), device="cuda")
|
||||
torch.cuda.synchronize()
|
||||
tic = time()
|
||||
with torch.inference_mode():
|
||||
audio_original = generator_original(data)
|
||||
torch.cuda.synchronize()
|
||||
toc = time() - tic
|
||||
toc_total_original += toc
|
||||
|
||||
vram_used_original_total += torch.cuda.max_memory_allocated(device="cuda")
|
||||
|
||||
del data, audio_original
|
||||
torch.cuda.empty_cache()
|
||||
|
||||
# Measure CUDA kernel inference in isolation
|
||||
for i in tqdm(range(num_sample)):
|
||||
torch.cuda.reset_peak_memory_stats(device="cuda")
|
||||
data = torch.rand((1, h.num_mels, num_mel_frame), device="cuda")
|
||||
torch.cuda.synchronize()
|
||||
tic = time()
|
||||
with torch.inference_mode():
|
||||
audio_cuda_kernel = generator_cuda_kernel(data)
|
||||
torch.cuda.synchronize()
|
||||
toc = time() - tic
|
||||
toc_total_cuda_kernel += toc
|
||||
|
||||
audio_length_total += audio_cuda_kernel.shape[-1]
|
||||
|
||||
vram_used_cuda_kernel_total += torch.cuda.max_memory_allocated(device="cuda")
|
||||
|
||||
del data, audio_cuda_kernel
|
||||
torch.cuda.empty_cache()
|
||||
|
||||
# Calculate metrics
|
||||
audio_second = audio_length_total / h.sampling_rate
|
||||
khz_original = audio_length_total / toc_total_original / 1000
|
||||
khz_cuda_kernel = audio_length_total / toc_total_cuda_kernel / 1000
|
||||
vram_used_original_gb = vram_used_original_total / num_sample / (1024**3)
|
||||
vram_used_cuda_kernel_gb = vram_used_cuda_kernel_total / num_sample / (1024**3)
|
||||
|
||||
# Print results
|
||||
print(
|
||||
f"Original BigVGAN: took {toc_total_original:.2f} seconds to generate {audio_second:.2f} seconds of audio, {khz_original:.1f}kHz, {audio_second / toc_total_original:.1f} faster than realtime, VRAM used {vram_used_original_gb:.1f} GB"
|
||||
)
|
||||
print(
|
||||
f"CUDA kernel BigVGAN: took {toc_total_cuda_kernel:.2f} seconds to generate {audio_second:.2f} seconds of audio, {khz_cuda_kernel:.1f}kHz, {audio_second / toc_total_cuda_kernel:.1f} faster than realtime, VRAM used {vram_used_cuda_kernel_gb:.1f} GB"
|
||||
)
|
||||
print(f"speedup of CUDA kernel: {khz_cuda_kernel / khz_original}")
|
||||
print(f"VRAM saving of CUDA kernel: {vram_used_original_gb / vram_used_cuda_kernel_gb}")
|
||||
|
||||
# Use artificial sine waves for inference test
|
||||
audio_real, sr = generate_soundwave(duration=5.0, sr=h.sampling_rate)
|
||||
audio_real = torch.tensor(audio_real).to("cuda")
|
||||
# Compute mel spectrogram from the ground truth audio
|
||||
x = get_mel(audio_real.unsqueeze(0), h)
|
||||
|
||||
with torch.inference_mode():
|
||||
y_g_hat_original = generator_original(x)
|
||||
y_g_hat_cuda_kernel = generator_cuda_kernel(x)
|
||||
|
||||
audio_real = audio_real.squeeze()
|
||||
audio_real = audio_real * MAX_WAV_VALUE
|
||||
audio_real = audio_real.cpu().numpy().astype("int16")
|
||||
|
||||
audio_original = y_g_hat_original.squeeze()
|
||||
audio_original = audio_original * MAX_WAV_VALUE
|
||||
audio_original = audio_original.cpu().numpy().astype("int16")
|
||||
|
||||
audio_cuda_kernel = y_g_hat_cuda_kernel.squeeze()
|
||||
audio_cuda_kernel = audio_cuda_kernel * MAX_WAV_VALUE
|
||||
audio_cuda_kernel = audio_cuda_kernel.cpu().numpy().astype("int16")
|
||||
|
||||
os.makedirs("tmp", exist_ok=True)
|
||||
output_file_real = os.path.join("tmp", "audio_real.wav")
|
||||
output_file_original = os.path.join("tmp", "audio_generated_original.wav")
|
||||
output_file_cuda_kernel = os.path.join("tmp", "audio_generated_cuda_kernel.wav")
|
||||
write(output_file_real, h.sampling_rate, audio_real)
|
||||
write(output_file_original, h.sampling_rate, audio_original)
|
||||
write(output_file_cuda_kernel, h.sampling_rate, audio_cuda_kernel)
|
||||
print("Example generated audios of original vs. fused CUDA kernel written to tmp!")
|
||||
print("Done")
|
||||
@@ -1,716 +0,0 @@
|
||||
# Copyright (c) 2024 NVIDIA CORPORATION.
|
||||
# Licensed under the MIT license.
|
||||
|
||||
# Adapted from https://github.com/jik876/hifi-gan under the MIT license.
|
||||
# LICENSE is in incl_licenses directory.
|
||||
|
||||
|
||||
import warnings
|
||||
|
||||
warnings.simplefilter(action="ignore", category=FutureWarning)
|
||||
import itertools
|
||||
import os
|
||||
import time
|
||||
import argparse
|
||||
import json
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
from torch.utils.tensorboard import SummaryWriter
|
||||
from torch.utils.data import DistributedSampler, DataLoader
|
||||
import torch.multiprocessing as mp
|
||||
from torch.distributed import init_process_group
|
||||
from torch.nn.parallel import DistributedDataParallel
|
||||
from env import AttrDict, build_env
|
||||
from meldataset import MelDataset, mel_spectrogram, get_dataset_filelist, MAX_WAV_VALUE
|
||||
|
||||
from bigvgan import BigVGAN
|
||||
from discriminators import (
|
||||
MultiPeriodDiscriminator,
|
||||
MultiResolutionDiscriminator,
|
||||
MultiBandDiscriminator,
|
||||
MultiScaleSubbandCQTDiscriminator,
|
||||
)
|
||||
from loss import (
|
||||
feature_loss,
|
||||
generator_loss,
|
||||
discriminator_loss,
|
||||
MultiScaleMelSpectrogramLoss,
|
||||
)
|
||||
|
||||
from utils import (
|
||||
plot_spectrogram,
|
||||
plot_spectrogram_clipped,
|
||||
scan_checkpoint,
|
||||
load_checkpoint,
|
||||
save_checkpoint,
|
||||
save_audio,
|
||||
)
|
||||
import torchaudio as ta
|
||||
from pesq import pesq
|
||||
from tqdm import tqdm
|
||||
import auraloss
|
||||
|
||||
torch.backends.cudnn.benchmark = False
|
||||
|
||||
|
||||
def train(rank, a, h):
|
||||
if h.num_gpus > 1:
|
||||
# initialize distributed
|
||||
init_process_group(
|
||||
backend=h.dist_config["dist_backend"],
|
||||
init_method=h.dist_config["dist_url"],
|
||||
world_size=h.dist_config["world_size"] * h.num_gpus,
|
||||
rank=rank,
|
||||
)
|
||||
|
||||
# Set seed and device
|
||||
torch.cuda.manual_seed(h.seed)
|
||||
torch.cuda.set_device(rank)
|
||||
device = torch.device(f"cuda:{rank:d}")
|
||||
|
||||
# Define BigVGAN generator
|
||||
generator = BigVGAN(h).to(device)
|
||||
|
||||
# Define discriminators. MPD is used by default
|
||||
mpd = MultiPeriodDiscriminator(h).to(device)
|
||||
|
||||
# Define additional discriminators. BigVGAN-v1 uses UnivNet's MRD as default
|
||||
# New in BigVGAN-v2: option to switch to new discriminators: MultiBandDiscriminator / MultiScaleSubbandCQTDiscriminator
|
||||
if h.get("use_mbd_instead_of_mrd", False): # Switch to MBD
|
||||
print("[INFO] using MultiBandDiscriminator of BigVGAN-v2 instead of MultiResolutionDiscriminator")
|
||||
# Variable name is kept as "mrd" for backward compatibility & minimal code change
|
||||
mrd = MultiBandDiscriminator(h).to(device)
|
||||
elif h.get("use_cqtd_instead_of_mrd", False): # Switch to CQTD
|
||||
print("[INFO] using MultiScaleSubbandCQTDiscriminator of BigVGAN-v2 instead of MultiResolutionDiscriminator")
|
||||
mrd = MultiScaleSubbandCQTDiscriminator(h).to(device)
|
||||
else: # Fallback to original MRD in BigVGAN-v1
|
||||
mrd = MultiResolutionDiscriminator(h).to(device)
|
||||
|
||||
# New in BigVGAN-v2: option to switch to multi-scale L1 mel loss
|
||||
if h.get("use_multiscale_melloss", False):
|
||||
print("[INFO] using multi-scale Mel l1 loss of BigVGAN-v2 instead of the original single-scale loss")
|
||||
fn_mel_loss_multiscale = MultiScaleMelSpectrogramLoss(
|
||||
sampling_rate=h.sampling_rate
|
||||
) # NOTE: accepts waveform as input
|
||||
else:
|
||||
fn_mel_loss_singlescale = F.l1_loss
|
||||
|
||||
# Print the model & number of parameters, and create or scan the latest checkpoint from checkpoints directory
|
||||
if rank == 0:
|
||||
print(generator)
|
||||
print(mpd)
|
||||
print(mrd)
|
||||
print(f"Generator params: {sum(p.numel() for p in generator.parameters())}")
|
||||
print(f"Discriminator mpd params: {sum(p.numel() for p in mpd.parameters())}")
|
||||
print(f"Discriminator mrd params: {sum(p.numel() for p in mrd.parameters())}")
|
||||
os.makedirs(a.checkpoint_path, exist_ok=True)
|
||||
print(f"Checkpoints directory: {a.checkpoint_path}")
|
||||
|
||||
if os.path.isdir(a.checkpoint_path):
|
||||
# New in v2.1: If the step prefix pattern-based checkpoints are not found, also check for renamed files in Hugging Face Hub to resume training
|
||||
cp_g = scan_checkpoint(a.checkpoint_path, prefix="g_", renamed_file="bigvgan_generator.pt")
|
||||
cp_do = scan_checkpoint(
|
||||
a.checkpoint_path,
|
||||
prefix="do_",
|
||||
renamed_file="bigvgan_discriminator_optimizer.pt",
|
||||
)
|
||||
|
||||
# Load the latest checkpoint if exists
|
||||
steps = 0
|
||||
if cp_g is None or cp_do is None:
|
||||
state_dict_do = None
|
||||
last_epoch = -1
|
||||
else:
|
||||
state_dict_g = load_checkpoint(cp_g, device)
|
||||
state_dict_do = load_checkpoint(cp_do, device)
|
||||
generator.load_state_dict(state_dict_g["generator"])
|
||||
mpd.load_state_dict(state_dict_do["mpd"])
|
||||
mrd.load_state_dict(state_dict_do["mrd"])
|
||||
steps = state_dict_do["steps"] + 1
|
||||
last_epoch = state_dict_do["epoch"]
|
||||
|
||||
# Initialize DDP, optimizers, and schedulers
|
||||
if h.num_gpus > 1:
|
||||
generator = DistributedDataParallel(generator, device_ids=[rank]).to(device)
|
||||
mpd = DistributedDataParallel(mpd, device_ids=[rank]).to(device)
|
||||
mrd = DistributedDataParallel(mrd, device_ids=[rank]).to(device)
|
||||
|
||||
optim_g = torch.optim.AdamW(generator.parameters(), h.learning_rate, betas=[h.adam_b1, h.adam_b2])
|
||||
optim_d = torch.optim.AdamW(
|
||||
itertools.chain(mrd.parameters(), mpd.parameters()),
|
||||
h.learning_rate,
|
||||
betas=[h.adam_b1, h.adam_b2],
|
||||
)
|
||||
|
||||
if state_dict_do is not None:
|
||||
optim_g.load_state_dict(state_dict_do["optim_g"])
|
||||
optim_d.load_state_dict(state_dict_do["optim_d"])
|
||||
|
||||
scheduler_g = torch.optim.lr_scheduler.ExponentialLR(optim_g, gamma=h.lr_decay, last_epoch=last_epoch)
|
||||
scheduler_d = torch.optim.lr_scheduler.ExponentialLR(optim_d, gamma=h.lr_decay, last_epoch=last_epoch)
|
||||
|
||||
# Define training and validation datasets
|
||||
|
||||
"""
|
||||
unseen_validation_filelist will contain sample filepaths outside the seen training & validation dataset
|
||||
Example: trained on LibriTTS, validate on VCTK
|
||||
"""
|
||||
training_filelist, validation_filelist, list_unseen_validation_filelist = get_dataset_filelist(a)
|
||||
|
||||
trainset = MelDataset(
|
||||
training_filelist,
|
||||
h,
|
||||
h.segment_size,
|
||||
h.n_fft,
|
||||
h.num_mels,
|
||||
h.hop_size,
|
||||
h.win_size,
|
||||
h.sampling_rate,
|
||||
h.fmin,
|
||||
h.fmax,
|
||||
shuffle=False if h.num_gpus > 1 else True,
|
||||
fmax_loss=h.fmax_for_loss,
|
||||
device=device,
|
||||
fine_tuning=a.fine_tuning,
|
||||
base_mels_path=a.input_mels_dir,
|
||||
is_seen=True,
|
||||
)
|
||||
|
||||
train_sampler = DistributedSampler(trainset) if h.num_gpus > 1 else None
|
||||
|
||||
train_loader = DataLoader(
|
||||
trainset,
|
||||
num_workers=h.num_workers,
|
||||
shuffle=False,
|
||||
sampler=train_sampler,
|
||||
batch_size=h.batch_size,
|
||||
pin_memory=True,
|
||||
drop_last=True,
|
||||
)
|
||||
|
||||
if rank == 0:
|
||||
validset = MelDataset(
|
||||
validation_filelist,
|
||||
h,
|
||||
h.segment_size,
|
||||
h.n_fft,
|
||||
h.num_mels,
|
||||
h.hop_size,
|
||||
h.win_size,
|
||||
h.sampling_rate,
|
||||
h.fmin,
|
||||
h.fmax,
|
||||
False,
|
||||
False,
|
||||
fmax_loss=h.fmax_for_loss,
|
||||
device=device,
|
||||
fine_tuning=a.fine_tuning,
|
||||
base_mels_path=a.input_mels_dir,
|
||||
is_seen=True,
|
||||
)
|
||||
validation_loader = DataLoader(
|
||||
validset,
|
||||
num_workers=1,
|
||||
shuffle=False,
|
||||
sampler=None,
|
||||
batch_size=1,
|
||||
pin_memory=True,
|
||||
drop_last=True,
|
||||
)
|
||||
|
||||
list_unseen_validset = []
|
||||
list_unseen_validation_loader = []
|
||||
for i in range(len(list_unseen_validation_filelist)):
|
||||
unseen_validset = MelDataset(
|
||||
list_unseen_validation_filelist[i],
|
||||
h,
|
||||
h.segment_size,
|
||||
h.n_fft,
|
||||
h.num_mels,
|
||||
h.hop_size,
|
||||
h.win_size,
|
||||
h.sampling_rate,
|
||||
h.fmin,
|
||||
h.fmax,
|
||||
False,
|
||||
False,
|
||||
fmax_loss=h.fmax_for_loss,
|
||||
device=device,
|
||||
fine_tuning=a.fine_tuning,
|
||||
base_mels_path=a.input_mels_dir,
|
||||
is_seen=False,
|
||||
)
|
||||
unseen_validation_loader = DataLoader(
|
||||
unseen_validset,
|
||||
num_workers=1,
|
||||
shuffle=False,
|
||||
sampler=None,
|
||||
batch_size=1,
|
||||
pin_memory=True,
|
||||
drop_last=True,
|
||||
)
|
||||
list_unseen_validset.append(unseen_validset)
|
||||
list_unseen_validation_loader.append(unseen_validation_loader)
|
||||
|
||||
# Tensorboard logger
|
||||
sw = SummaryWriter(os.path.join(a.checkpoint_path, "logs"))
|
||||
if a.save_audio: # Also save audio to disk if --save_audio is set to True
|
||||
os.makedirs(os.path.join(a.checkpoint_path, "samples"), exist_ok=True)
|
||||
|
||||
"""
|
||||
Validation loop, "mode" parameter is automatically defined as (seen or unseen)_(name of the dataset).
|
||||
If the name of the dataset contains "nonspeech", it skips PESQ calculation to prevent errors
|
||||
"""
|
||||
|
||||
def validate(rank, a, h, loader, mode="seen"):
|
||||
assert rank == 0, "validate should only run on rank=0"
|
||||
generator.eval()
|
||||
torch.cuda.empty_cache()
|
||||
|
||||
val_err_tot = 0
|
||||
val_pesq_tot = 0
|
||||
val_mrstft_tot = 0
|
||||
|
||||
# Modules for evaluation metrics
|
||||
pesq_resampler = ta.transforms.Resample(h.sampling_rate, 16000).cuda()
|
||||
loss_mrstft = auraloss.freq.MultiResolutionSTFTLoss(device="cuda")
|
||||
|
||||
if a.save_audio: # Also save audio to disk if --save_audio is set to True
|
||||
os.makedirs(
|
||||
os.path.join(a.checkpoint_path, "samples", f"gt_{mode}"),
|
||||
exist_ok=True,
|
||||
)
|
||||
os.makedirs(
|
||||
os.path.join(a.checkpoint_path, "samples", f"{mode}_{steps:08d}"),
|
||||
exist_ok=True,
|
||||
)
|
||||
|
||||
with torch.no_grad():
|
||||
print(f"step {steps} {mode} speaker validation...")
|
||||
|
||||
# Loop over validation set and compute metrics
|
||||
for j, batch in enumerate(tqdm(loader)):
|
||||
x, y, _, y_mel = batch
|
||||
y = y.to(device)
|
||||
if hasattr(generator, "module"):
|
||||
y_g_hat = generator.module(x.to(device))
|
||||
else:
|
||||
y_g_hat = generator(x.to(device))
|
||||
y_mel = y_mel.to(device, non_blocking=True)
|
||||
y_g_hat_mel = mel_spectrogram(
|
||||
y_g_hat.squeeze(1),
|
||||
h.n_fft,
|
||||
h.num_mels,
|
||||
h.sampling_rate,
|
||||
h.hop_size,
|
||||
h.win_size,
|
||||
h.fmin,
|
||||
h.fmax_for_loss,
|
||||
)
|
||||
min_t = min(y_mel.size(-1), y_g_hat_mel.size(-1))
|
||||
val_err_tot += F.l1_loss(y_mel[..., :min_t], y_g_hat_mel[..., :min_t]).item()
|
||||
|
||||
# PESQ calculation. only evaluate PESQ if it's speech signal (nonspeech PESQ will error out)
|
||||
if "nonspeech" not in mode: # Skips if the name of dataset (in mode string) contains "nonspeech"
|
||||
# Resample to 16000 for pesq
|
||||
y_16k = pesq_resampler(y)
|
||||
y_g_hat_16k = pesq_resampler(y_g_hat.squeeze(1))
|
||||
y_int_16k = (y_16k[0] * MAX_WAV_VALUE).short().cpu().numpy()
|
||||
y_g_hat_int_16k = (y_g_hat_16k[0] * MAX_WAV_VALUE).short().cpu().numpy()
|
||||
val_pesq_tot += pesq(16000, y_int_16k, y_g_hat_int_16k, "wb")
|
||||
|
||||
# MRSTFT calculation
|
||||
min_t = min(y.size(-1), y_g_hat.size(-1))
|
||||
val_mrstft_tot += loss_mrstft(y_g_hat[..., :min_t], y[..., :min_t]).item()
|
||||
|
||||
# Log audio and figures to Tensorboard
|
||||
if j % a.eval_subsample == 0: # Subsample every nth from validation set
|
||||
if steps >= 0:
|
||||
sw.add_audio(f"gt_{mode}/y_{j}", y[0], steps, h.sampling_rate)
|
||||
if a.save_audio: # Also save audio to disk if --save_audio is set to True
|
||||
save_audio(
|
||||
y[0],
|
||||
os.path.join(
|
||||
a.checkpoint_path,
|
||||
"samples",
|
||||
f"gt_{mode}",
|
||||
f"{j:04d}.wav",
|
||||
),
|
||||
h.sampling_rate,
|
||||
)
|
||||
sw.add_figure(
|
||||
f"gt_{mode}/y_spec_{j}",
|
||||
plot_spectrogram(x[0]),
|
||||
steps,
|
||||
)
|
||||
|
||||
sw.add_audio(
|
||||
f"generated_{mode}/y_hat_{j}",
|
||||
y_g_hat[0],
|
||||
steps,
|
||||
h.sampling_rate,
|
||||
)
|
||||
if a.save_audio: # Also save audio to disk if --save_audio is set to True
|
||||
save_audio(
|
||||
y_g_hat[0, 0],
|
||||
os.path.join(
|
||||
a.checkpoint_path,
|
||||
"samples",
|
||||
f"{mode}_{steps:08d}",
|
||||
f"{j:04d}.wav",
|
||||
),
|
||||
h.sampling_rate,
|
||||
)
|
||||
# Spectrogram of synthesized audio
|
||||
y_hat_spec = mel_spectrogram(
|
||||
y_g_hat.squeeze(1),
|
||||
h.n_fft,
|
||||
h.num_mels,
|
||||
h.sampling_rate,
|
||||
h.hop_size,
|
||||
h.win_size,
|
||||
h.fmin,
|
||||
h.fmax,
|
||||
)
|
||||
sw.add_figure(
|
||||
f"generated_{mode}/y_hat_spec_{j}",
|
||||
plot_spectrogram(y_hat_spec.squeeze(0).cpu().numpy()),
|
||||
steps,
|
||||
)
|
||||
|
||||
"""
|
||||
Visualization of spectrogram difference between GT and synthesized audio, difference higher than 1 is clipped for better visualization.
|
||||
"""
|
||||
spec_delta = torch.clamp(
|
||||
torch.abs(x[0] - y_hat_spec.squeeze(0).cpu()),
|
||||
min=1e-6,
|
||||
max=1.0,
|
||||
)
|
||||
sw.add_figure(
|
||||
f"delta_dclip1_{mode}/spec_{j}",
|
||||
plot_spectrogram_clipped(spec_delta.numpy(), clip_max=1.0),
|
||||
steps,
|
||||
)
|
||||
|
||||
val_err = val_err_tot / (j + 1)
|
||||
val_pesq = val_pesq_tot / (j + 1)
|
||||
val_mrstft = val_mrstft_tot / (j + 1)
|
||||
# Log evaluation metrics to Tensorboard
|
||||
sw.add_scalar(f"validation_{mode}/mel_spec_error", val_err, steps)
|
||||
sw.add_scalar(f"validation_{mode}/pesq", val_pesq, steps)
|
||||
sw.add_scalar(f"validation_{mode}/mrstft", val_mrstft, steps)
|
||||
|
||||
generator.train()
|
||||
|
||||
# If the checkpoint is loaded, start with validation loop
|
||||
if steps != 0 and rank == 0 and not a.debug:
|
||||
if not a.skip_seen:
|
||||
validate(
|
||||
rank,
|
||||
a,
|
||||
h,
|
||||
validation_loader,
|
||||
mode=f"seen_{train_loader.dataset.name}",
|
||||
)
|
||||
for i in range(len(list_unseen_validation_loader)):
|
||||
validate(
|
||||
rank,
|
||||
a,
|
||||
h,
|
||||
list_unseen_validation_loader[i],
|
||||
mode=f"unseen_{list_unseen_validation_loader[i].dataset.name}",
|
||||
)
|
||||
# Exit the script if --evaluate is set to True
|
||||
if a.evaluate:
|
||||
exit()
|
||||
|
||||
# Main training loop
|
||||
generator.train()
|
||||
mpd.train()
|
||||
mrd.train()
|
||||
for epoch in range(max(0, last_epoch), a.training_epochs):
|
||||
if rank == 0:
|
||||
start = time.time()
|
||||
print(f"Epoch: {epoch + 1}")
|
||||
|
||||
if h.num_gpus > 1:
|
||||
train_sampler.set_epoch(epoch)
|
||||
|
||||
for i, batch in enumerate(train_loader):
|
||||
if rank == 0:
|
||||
start_b = time.time()
|
||||
x, y, _, y_mel = batch
|
||||
|
||||
x = x.to(device, non_blocking=True)
|
||||
y = y.to(device, non_blocking=True)
|
||||
y_mel = y_mel.to(device, non_blocking=True)
|
||||
y = y.unsqueeze(1)
|
||||
|
||||
y_g_hat = generator(x)
|
||||
y_g_hat_mel = mel_spectrogram(
|
||||
y_g_hat.squeeze(1),
|
||||
h.n_fft,
|
||||
h.num_mels,
|
||||
h.sampling_rate,
|
||||
h.hop_size,
|
||||
h.win_size,
|
||||
h.fmin,
|
||||
h.fmax_for_loss,
|
||||
)
|
||||
|
||||
optim_d.zero_grad()
|
||||
|
||||
# MPD
|
||||
y_df_hat_r, y_df_hat_g, _, _ = mpd(y, y_g_hat.detach())
|
||||
loss_disc_f, losses_disc_f_r, losses_disc_f_g = discriminator_loss(y_df_hat_r, y_df_hat_g)
|
||||
|
||||
# MRD
|
||||
y_ds_hat_r, y_ds_hat_g, _, _ = mrd(y, y_g_hat.detach())
|
||||
loss_disc_s, losses_disc_s_r, losses_disc_s_g = discriminator_loss(y_ds_hat_r, y_ds_hat_g)
|
||||
|
||||
loss_disc_all = loss_disc_s + loss_disc_f
|
||||
|
||||
# Set clip_grad_norm value
|
||||
clip_grad_norm = h.get("clip_grad_norm", 1000.0) # Default to 1000
|
||||
|
||||
# Whether to freeze D for initial training steps
|
||||
if steps >= a.freeze_step:
|
||||
loss_disc_all.backward()
|
||||
grad_norm_mpd = torch.nn.utils.clip_grad_norm_(mpd.parameters(), clip_grad_norm)
|
||||
grad_norm_mrd = torch.nn.utils.clip_grad_norm_(mrd.parameters(), clip_grad_norm)
|
||||
optim_d.step()
|
||||
else:
|
||||
print(f"[WARNING] skipping D training for the first {a.freeze_step} steps")
|
||||
grad_norm_mpd = 0.0
|
||||
grad_norm_mrd = 0.0
|
||||
|
||||
# Generator
|
||||
optim_g.zero_grad()
|
||||
|
||||
# L1 Mel-Spectrogram Loss
|
||||
lambda_melloss = h.get("lambda_melloss", 45.0) # Defaults to 45 in BigVGAN-v1 if not set
|
||||
if h.get("use_multiscale_melloss", False): # uses wav <y, y_g_hat> for loss
|
||||
loss_mel = fn_mel_loss_multiscale(y, y_g_hat) * lambda_melloss
|
||||
else: # Uses mel <y_mel, y_g_hat_mel> for loss
|
||||
loss_mel = fn_mel_loss_singlescale(y_mel, y_g_hat_mel) * lambda_melloss
|
||||
|
||||
# MPD loss
|
||||
y_df_hat_r, y_df_hat_g, fmap_f_r, fmap_f_g = mpd(y, y_g_hat)
|
||||
loss_fm_f = feature_loss(fmap_f_r, fmap_f_g)
|
||||
loss_gen_f, losses_gen_f = generator_loss(y_df_hat_g)
|
||||
|
||||
# MRD loss
|
||||
y_ds_hat_r, y_ds_hat_g, fmap_s_r, fmap_s_g = mrd(y, y_g_hat)
|
||||
loss_fm_s = feature_loss(fmap_s_r, fmap_s_g)
|
||||
loss_gen_s, losses_gen_s = generator_loss(y_ds_hat_g)
|
||||
|
||||
if steps >= a.freeze_step:
|
||||
loss_gen_all = loss_gen_s + loss_gen_f + loss_fm_s + loss_fm_f + loss_mel
|
||||
else:
|
||||
print(f"[WARNING] using regression loss only for G for the first {a.freeze_step} steps")
|
||||
loss_gen_all = loss_mel
|
||||
|
||||
loss_gen_all.backward()
|
||||
grad_norm_g = torch.nn.utils.clip_grad_norm_(generator.parameters(), clip_grad_norm)
|
||||
optim_g.step()
|
||||
|
||||
if rank == 0:
|
||||
# STDOUT logging
|
||||
if steps % a.stdout_interval == 0:
|
||||
mel_error = loss_mel.item() / lambda_melloss # Log training mel regression loss to stdout
|
||||
print(
|
||||
f"Steps: {steps:d}, "
|
||||
f"Gen Loss Total: {loss_gen_all:4.3f}, "
|
||||
f"Mel Error: {mel_error:4.3f}, "
|
||||
f"s/b: {time.time() - start_b:4.3f} "
|
||||
f"lr: {optim_g.param_groups[0]['lr']:4.7f} "
|
||||
f"grad_norm_g: {grad_norm_g:4.3f}"
|
||||
)
|
||||
|
||||
# Checkpointing
|
||||
if steps % a.checkpoint_interval == 0 and steps != 0:
|
||||
checkpoint_path = f"{a.checkpoint_path}/g_{steps:08d}"
|
||||
save_checkpoint(
|
||||
checkpoint_path,
|
||||
{"generator": (generator.module if h.num_gpus > 1 else generator).state_dict()},
|
||||
)
|
||||
checkpoint_path = f"{a.checkpoint_path}/do_{steps:08d}"
|
||||
save_checkpoint(
|
||||
checkpoint_path,
|
||||
{
|
||||
"mpd": (mpd.module if h.num_gpus > 1 else mpd).state_dict(),
|
||||
"mrd": (mrd.module if h.num_gpus > 1 else mrd).state_dict(),
|
||||
"optim_g": optim_g.state_dict(),
|
||||
"optim_d": optim_d.state_dict(),
|
||||
"steps": steps,
|
||||
"epoch": epoch,
|
||||
},
|
||||
)
|
||||
|
||||
# Tensorboard summary logging
|
||||
if steps % a.summary_interval == 0:
|
||||
mel_error = loss_mel.item() / lambda_melloss # Log training mel regression loss to tensorboard
|
||||
sw.add_scalar("training/gen_loss_total", loss_gen_all.item(), steps)
|
||||
sw.add_scalar("training/mel_spec_error", mel_error, steps)
|
||||
sw.add_scalar("training/fm_loss_mpd", loss_fm_f.item(), steps)
|
||||
sw.add_scalar("training/gen_loss_mpd", loss_gen_f.item(), steps)
|
||||
sw.add_scalar("training/disc_loss_mpd", loss_disc_f.item(), steps)
|
||||
sw.add_scalar("training/grad_norm_mpd", grad_norm_mpd, steps)
|
||||
sw.add_scalar("training/fm_loss_mrd", loss_fm_s.item(), steps)
|
||||
sw.add_scalar("training/gen_loss_mrd", loss_gen_s.item(), steps)
|
||||
sw.add_scalar("training/disc_loss_mrd", loss_disc_s.item(), steps)
|
||||
sw.add_scalar("training/grad_norm_mrd", grad_norm_mrd, steps)
|
||||
sw.add_scalar("training/grad_norm_g", grad_norm_g, steps)
|
||||
sw.add_scalar("training/learning_rate_d", scheduler_d.get_last_lr()[0], steps)
|
||||
sw.add_scalar("training/learning_rate_g", scheduler_g.get_last_lr()[0], steps)
|
||||
sw.add_scalar("training/epoch", epoch + 1, steps)
|
||||
|
||||
# Validation
|
||||
if steps % a.validation_interval == 0:
|
||||
# Plot training input x so far used
|
||||
for i_x in range(x.shape[0]):
|
||||
sw.add_figure(
|
||||
f"training_input/x_{i_x}",
|
||||
plot_spectrogram(x[i_x].cpu()),
|
||||
steps,
|
||||
)
|
||||
sw.add_audio(
|
||||
f"training_input/y_{i_x}",
|
||||
y[i_x][0],
|
||||
steps,
|
||||
h.sampling_rate,
|
||||
)
|
||||
|
||||
# Seen and unseen speakers validation loops
|
||||
if not a.debug and steps != 0:
|
||||
validate(
|
||||
rank,
|
||||
a,
|
||||
h,
|
||||
validation_loader,
|
||||
mode=f"seen_{train_loader.dataset.name}",
|
||||
)
|
||||
for i in range(len(list_unseen_validation_loader)):
|
||||
validate(
|
||||
rank,
|
||||
a,
|
||||
h,
|
||||
list_unseen_validation_loader[i],
|
||||
mode=f"unseen_{list_unseen_validation_loader[i].dataset.name}",
|
||||
)
|
||||
steps += 1
|
||||
|
||||
# BigVGAN-v2 learning rate scheduler is changed from epoch-level to step-level
|
||||
scheduler_g.step()
|
||||
scheduler_d.step()
|
||||
|
||||
if rank == 0:
|
||||
print(f"Time taken for epoch {epoch + 1} is {int(time.time() - start)} sec\n")
|
||||
|
||||
|
||||
def main():
|
||||
print("Initializing Training Process..")
|
||||
|
||||
parser = argparse.ArgumentParser()
|
||||
|
||||
parser.add_argument("--group_name", default=None)
|
||||
|
||||
parser.add_argument("--input_wavs_dir", default="LibriTTS")
|
||||
parser.add_argument("--input_mels_dir", default="ft_dataset")
|
||||
parser.add_argument("--input_training_file", default="tests/LibriTTS/train-full.txt")
|
||||
parser.add_argument("--input_validation_file", default="tests/LibriTTS/val-full.txt")
|
||||
|
||||
parser.add_argument(
|
||||
"--list_input_unseen_wavs_dir",
|
||||
nargs="+",
|
||||
default=["tests/LibriTTS", "tests/LibriTTS"],
|
||||
)
|
||||
parser.add_argument(
|
||||
"--list_input_unseen_validation_file",
|
||||
nargs="+",
|
||||
default=["tests/LibriTTS/dev-clean.txt", "tests/LibriTTS/dev-other.txt"],
|
||||
)
|
||||
|
||||
parser.add_argument("--checkpoint_path", default="exp/bigvgan")
|
||||
parser.add_argument("--config", default="")
|
||||
|
||||
parser.add_argument("--training_epochs", default=100000, type=int)
|
||||
parser.add_argument("--stdout_interval", default=5, type=int)
|
||||
parser.add_argument("--checkpoint_interval", default=50000, type=int)
|
||||
parser.add_argument("--summary_interval", default=100, type=int)
|
||||
parser.add_argument("--validation_interval", default=50000, type=int)
|
||||
|
||||
parser.add_argument(
|
||||
"--freeze_step",
|
||||
default=0,
|
||||
type=int,
|
||||
help="freeze D for the first specified steps. G only uses regression loss for these steps.",
|
||||
)
|
||||
|
||||
parser.add_argument("--fine_tuning", default=False, type=bool)
|
||||
|
||||
parser.add_argument(
|
||||
"--debug",
|
||||
default=False,
|
||||
type=bool,
|
||||
help="debug mode. skips validation loop throughout training",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--evaluate",
|
||||
default=False,
|
||||
type=bool,
|
||||
help="only run evaluation from checkpoint and exit",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--eval_subsample",
|
||||
default=5,
|
||||
type=int,
|
||||
help="subsampling during evaluation loop",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--skip_seen",
|
||||
default=False,
|
||||
type=bool,
|
||||
help="skip seen dataset. useful for test set inference",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--save_audio",
|
||||
default=False,
|
||||
type=bool,
|
||||
help="save audio of test set inference to disk",
|
||||
)
|
||||
|
||||
a = parser.parse_args()
|
||||
|
||||
with open(a.config) as f:
|
||||
data = f.read()
|
||||
|
||||
json_config = json.loads(data)
|
||||
h = AttrDict(json_config)
|
||||
|
||||
build_env(a.config, "config.json", a.checkpoint_path)
|
||||
|
||||
torch.manual_seed(h.seed)
|
||||
if torch.cuda.is_available():
|
||||
torch.cuda.manual_seed(h.seed)
|
||||
h.num_gpus = torch.cuda.device_count()
|
||||
h.batch_size = int(h.batch_size / h.num_gpus)
|
||||
print(f"Batch size per GPU: {h.batch_size}")
|
||||
else:
|
||||
pass
|
||||
|
||||
if h.num_gpus > 1:
|
||||
mp.spawn(
|
||||
train,
|
||||
nprocs=h.num_gpus,
|
||||
args=(
|
||||
a,
|
||||
h,
|
||||
),
|
||||
)
|
||||
else:
|
||||
train(0, a, h)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -1,45 +1,5 @@
|
||||
# Adapted from https://github.com/jik876/hifi-gan under the MIT license.
|
||||
# LICENSE is in incl_licenses directory.
|
||||
|
||||
import glob
|
||||
import os
|
||||
import matplotlib
|
||||
import torch
|
||||
from torch.nn.utils import weight_norm
|
||||
|
||||
matplotlib.use("Agg")
|
||||
import matplotlib.pylab as plt
|
||||
from .meldataset import MAX_WAV_VALUE
|
||||
from scipy.io.wavfile import write
|
||||
|
||||
|
||||
def plot_spectrogram(spectrogram):
|
||||
fig, ax = plt.subplots(figsize=(10, 2))
|
||||
im = ax.imshow(spectrogram, aspect="auto", origin="lower", interpolation="none")
|
||||
plt.colorbar(im, ax=ax)
|
||||
|
||||
fig.canvas.draw()
|
||||
plt.close()
|
||||
|
||||
return fig
|
||||
|
||||
|
||||
def plot_spectrogram_clipped(spectrogram, clip_max=2.0):
|
||||
fig, ax = plt.subplots(figsize=(10, 2))
|
||||
im = ax.imshow(
|
||||
spectrogram,
|
||||
aspect="auto",
|
||||
origin="lower",
|
||||
interpolation="none",
|
||||
vmin=1e-6,
|
||||
vmax=clip_max,
|
||||
)
|
||||
plt.colorbar(im, ax=ax)
|
||||
|
||||
fig.canvas.draw()
|
||||
plt.close()
|
||||
|
||||
return fig
|
||||
# LICENSE is in the incl_licenses directory.
|
||||
|
||||
|
||||
def init_weights(m, mean=0.0, std=0.01):
|
||||
@@ -48,52 +8,5 @@ def init_weights(m, mean=0.0, std=0.01):
|
||||
m.weight.data.normal_(mean, std)
|
||||
|
||||
|
||||
def apply_weight_norm(m):
|
||||
classname = m.__class__.__name__
|
||||
if classname.find("Conv") != -1:
|
||||
weight_norm(m)
|
||||
|
||||
|
||||
def get_padding(kernel_size, dilation=1):
|
||||
return int((kernel_size * dilation - dilation) / 2)
|
||||
|
||||
|
||||
def load_checkpoint(filepath, device):
|
||||
assert os.path.isfile(filepath)
|
||||
print(f"Loading '{filepath}'")
|
||||
checkpoint_dict = torch.load(filepath, map_location=device)
|
||||
print("Complete.")
|
||||
return checkpoint_dict
|
||||
|
||||
|
||||
def save_checkpoint(filepath, obj):
|
||||
print(f"Saving checkpoint to {filepath}")
|
||||
torch.save(obj, filepath)
|
||||
print("Complete.")
|
||||
|
||||
|
||||
def scan_checkpoint(cp_dir, prefix, renamed_file=None):
|
||||
# Fallback to original scanning logic first
|
||||
pattern = os.path.join(cp_dir, prefix + "????????")
|
||||
cp_list = glob.glob(pattern)
|
||||
|
||||
if len(cp_list) > 0:
|
||||
last_checkpoint_path = sorted(cp_list)[-1]
|
||||
print(f"[INFO] Resuming from checkpoint: '{last_checkpoint_path}'")
|
||||
return last_checkpoint_path
|
||||
|
||||
# If no pattern-based checkpoints are found, check for renamed file
|
||||
if renamed_file:
|
||||
renamed_path = os.path.join(cp_dir, renamed_file)
|
||||
if os.path.isfile(renamed_path):
|
||||
print(f"[INFO] Resuming from renamed checkpoint: '{renamed_file}'")
|
||||
return renamed_path
|
||||
|
||||
return None
|
||||
|
||||
|
||||
def save_audio(audio, path, sr):
|
||||
# wav: torch with 1d shape
|
||||
audio = audio * MAX_WAV_VALUE
|
||||
audio = audio.cpu().numpy().astype("int16")
|
||||
write(path, sr, audio)
|
||||
|
||||
+372
-156
@@ -32,19 +32,21 @@ from transformers import AutoModelForMaskedLM, AutoTokenizer
|
||||
|
||||
from tools.audio_sr import AP_BWE
|
||||
from tools.i18n.i18n import I18nAuto, scan_language_list
|
||||
from tools.my_utils import load_audio
|
||||
from TTS_infer_pack.text_segmentation_method import splits
|
||||
from TTS_infer_pack.TextPreprocessor import TextPreprocessor
|
||||
from sv import SV
|
||||
resample_transform_dict={}
|
||||
def resample(audio_tensor, sr0,sr1,device):
|
||||
|
||||
resample_transform_dict = {}
|
||||
|
||||
|
||||
def resample(audio_tensor, sr0, sr1, device):
|
||||
global resample_transform_dict
|
||||
key="%s-%s-%s"%(sr0,sr1,str(device))
|
||||
key = "%s-%s-%s" % (sr0, sr1, str(device))
|
||||
if key not in resample_transform_dict:
|
||||
resample_transform_dict[key] = torchaudio.transforms.Resample(
|
||||
sr0, sr1
|
||||
).to(device)
|
||||
resample_transform_dict[key] = torchaudio.transforms.Resample(sr0, sr1).to(device)
|
||||
return resample_transform_dict[key](audio_tensor)
|
||||
|
||||
|
||||
language = os.environ.get("language", "Auto")
|
||||
language = sys.argv[-1] if sys.argv[-1] in scan_language_list() else language
|
||||
i18n = I18nAuto(language=language)
|
||||
@@ -111,6 +113,7 @@ def speed_change(input_audio: np.ndarray, speed: float, sr: int):
|
||||
|
||||
return processed_audio
|
||||
|
||||
|
||||
class DictToAttrRecursive(dict):
|
||||
def __init__(self, input_dict):
|
||||
super().__init__(input_dict)
|
||||
@@ -272,6 +275,15 @@ class TTS_Config:
|
||||
v1_languages: list = ["auto", "en", "zh", "ja", "all_zh", "all_ja"]
|
||||
v2_languages: list = ["auto", "auto_yue", "en", "zh", "ja", "yue", "ko", "all_zh", "all_ja", "all_yue", "all_ko"]
|
||||
languages: list = v2_languages
|
||||
mute_tokens: dict = {
|
||||
"v1" : 486,
|
||||
"v2" : 486,
|
||||
"v2Pro": 486,
|
||||
"v2ProPlus": 486,
|
||||
"v3" : 486,
|
||||
"v4" : 486,
|
||||
}
|
||||
mute_emb_sim_matrix: torch.Tensor = None
|
||||
# "all_zh",#全部按中文识别
|
||||
# "en",#全部按英文识别#######不变
|
||||
# "all_ja",#全部按日文识别
|
||||
@@ -301,10 +313,10 @@ class TTS_Config:
|
||||
configs: dict = self._load_configs(self.configs_path)
|
||||
|
||||
assert isinstance(configs, dict)
|
||||
version = configs.get("version", "v2").lower()
|
||||
assert version in ["v1", "v2", "v3", "v4", "v2Pro", "v2ProPlus"]
|
||||
self.default_configs[version] = configs.get(version, self.default_configs[version])
|
||||
self.configs: dict = configs.get("custom", deepcopy(self.default_configs[version]))
|
||||
configs_ = deepcopy(self.default_configs)
|
||||
configs_.update(configs)
|
||||
self.configs: dict = configs_.get("custom", configs_["v2"])
|
||||
self.default_configs = deepcopy(configs_)
|
||||
|
||||
self.device = self.configs.get("device", torch.device("cpu"))
|
||||
if "cuda" in str(self.device) and not torch.cuda.is_available():
|
||||
@@ -312,11 +324,13 @@ class TTS_Config:
|
||||
self.device = torch.device("cpu")
|
||||
|
||||
self.is_half = self.configs.get("is_half", False)
|
||||
# if str(self.device) == "cpu" and self.is_half:
|
||||
# print(f"Warning: Half precision is not supported on CPU, set is_half to False.")
|
||||
# self.is_half = False
|
||||
if str(self.device) == "cpu" and self.is_half:
|
||||
print(f"Warning: Half precision is not supported on CPU, set is_half to False.")
|
||||
self.is_half = False
|
||||
|
||||
version = self.configs.get("version", None)
|
||||
self.version = version
|
||||
assert self.version in ["v1", "v2", "v3", "v4", "v2Pro", "v2ProPlus"], "Invalid version!"
|
||||
self.t2s_weights_path = self.configs.get("t2s_weights_path", None)
|
||||
self.vits_weights_path = self.configs.get("vits_weights_path", None)
|
||||
self.bert_base_path = self.configs.get("bert_base_path", None)
|
||||
@@ -479,13 +493,13 @@ class TTS:
|
||||
def init_vits_weights(self, weights_path: str):
|
||||
self.configs.vits_weights_path = weights_path
|
||||
version, model_version, if_lora_v3 = get_sovits_version_from_path_fast(weights_path)
|
||||
if "Pro"in model_version:
|
||||
if "Pro" in model_version:
|
||||
self.init_sv_model()
|
||||
path_sovits = self.configs.default_configs[model_version]["vits_weights_path"]
|
||||
|
||||
if if_lora_v3 == True and os.path.exists(path_sovits) == False:
|
||||
info = path_sovits + i18n("SoVITS %s 底模缺失,无法加载相应 LoRA 权重" % model_version)
|
||||
raise FileExistsError(info)
|
||||
raise FileNotFoundError(info)
|
||||
|
||||
# dict_s2 = torch.load(weights_path, map_location=self.configs.device,weights_only=False)
|
||||
dict_s2 = load_sovits_new(weights_path)
|
||||
@@ -498,9 +512,9 @@ class TTS:
|
||||
else:
|
||||
hps["model"]["version"] = "v2"
|
||||
version = hps["model"]["version"]
|
||||
v3v4set={"v3", "v4"}
|
||||
v3v4set = {"v3", "v4"}
|
||||
if model_version not in v3v4set:
|
||||
if "Pro"not in model_version:
|
||||
if "Pro" not in model_version:
|
||||
model_version = version
|
||||
else:
|
||||
hps["model"]["version"] = model_version
|
||||
@@ -542,7 +556,7 @@ class TTS:
|
||||
if "pretrained" not in weights_path and hasattr(vits_model, "enc_q"):
|
||||
del vits_model.enc_q
|
||||
|
||||
self.is_v2pro=model_version in {"v2Pro","v2ProPlus"}
|
||||
self.is_v2pro = model_version in {"v2Pro", "v2ProPlus"}
|
||||
|
||||
if if_lora_v3 == False:
|
||||
print(
|
||||
@@ -573,6 +587,10 @@ class TTS:
|
||||
if self.configs.is_half and str(self.configs.device) != "cpu":
|
||||
self.vits_model = self.vits_model.half()
|
||||
|
||||
self.configs.save_configs()
|
||||
|
||||
|
||||
|
||||
def init_t2s_weights(self, weights_path: str):
|
||||
print(f"Loading Text2Semantic weights from {weights_path}")
|
||||
self.configs.t2s_weights_path = weights_path
|
||||
@@ -589,6 +607,11 @@ class TTS:
|
||||
if self.configs.is_half and str(self.configs.device) != "cpu":
|
||||
self.t2s_model = self.t2s_model.half()
|
||||
|
||||
codebook = t2s_model.model.ar_audio_embedding.weight.clone()
|
||||
mute_emb = codebook[self.configs.mute_tokens[self.configs.version]].unsqueeze(0)
|
||||
sim_matrix = F.cosine_similarity(mute_emb.float(), codebook.float(), dim=-1)
|
||||
self.configs.mute_emb_sim_matrix = sim_matrix
|
||||
|
||||
def init_vocoder(self, version: str):
|
||||
if version == "v3":
|
||||
if self.vocoder is not None and self.vocoder.__class__.__name__ == "BigVGAN":
|
||||
@@ -600,8 +623,7 @@ class TTS:
|
||||
|
||||
self.vocoder = BigVGAN.from_pretrained(
|
||||
"%s/GPT_SoVITS/pretrained_models/models--nvidia--bigvgan_v2_24khz_100band_256x" % (now_dir,),
|
||||
use_cuda_kernel=False,
|
||||
) # if True, RuntimeError: Ninja is required to load C++ extensions
|
||||
)
|
||||
# remove weight norm in the model and set to eval mode
|
||||
self.vocoder.remove_weight_norm()
|
||||
|
||||
@@ -632,7 +654,9 @@ class TTS:
|
||||
)
|
||||
self.vocoder.remove_weight_norm()
|
||||
state_dict_g = torch.load(
|
||||
"%s/GPT_SoVITS/pretrained_models/gsv-v4-pretrained/vocoder.pth" % (now_dir,), map_location="cpu", weights_only=False
|
||||
"%s/GPT_SoVITS/pretrained_models/gsv-v4-pretrained/vocoder.pth" % (now_dir,),
|
||||
map_location="cpu",
|
||||
weights_only=False,
|
||||
)
|
||||
print("loading vocoder", self.vocoder.load_state_dict(state_dict_g))
|
||||
|
||||
@@ -752,11 +776,13 @@ class TTS:
|
||||
|
||||
if raw_sr != self.configs.sampling_rate:
|
||||
audio = raw_audio.to(self.configs.device)
|
||||
if (audio.shape[0] == 2): audio = audio.mean(0).unsqueeze(0)
|
||||
if audio.shape[0] == 2:
|
||||
audio = audio.mean(0).unsqueeze(0)
|
||||
audio = resample(audio, raw_sr, self.configs.sampling_rate, self.configs.device)
|
||||
else:
|
||||
audio = raw_audio.to(self.configs.device)
|
||||
if (audio.shape[0] == 2): audio = audio.mean(0).unsqueeze(0)
|
||||
if audio.shape[0] == 2:
|
||||
audio = audio.mean(0).unsqueeze(0)
|
||||
|
||||
maxx = audio.abs().max()
|
||||
if maxx > 1:
|
||||
@@ -775,8 +801,9 @@ class TTS:
|
||||
audio = resample(audio, self.configs.sampling_rate, 16000, self.configs.device)
|
||||
if self.configs.is_half:
|
||||
audio = audio.half()
|
||||
else:audio=None
|
||||
return spec,audio
|
||||
else:
|
||||
audio = None
|
||||
return spec, audio
|
||||
|
||||
def _set_prompt_semantic(self, ref_wav_path: str):
|
||||
zero_wav = np.zeros(
|
||||
@@ -980,21 +1007,25 @@ class TTS:
|
||||
"aux_ref_audio_paths": [], # list.(optional) auxiliary reference audio paths for multi-speaker tone fusion
|
||||
"prompt_text": "", # str.(optional) prompt text for the reference audio
|
||||
"prompt_lang": "", # str.(required) language of the prompt text for the reference audio
|
||||
"top_k": 5, # int. top k sampling
|
||||
"top_k": 15, # int. top k sampling
|
||||
"top_p": 1, # float. top p sampling
|
||||
"temperature": 1, # float. temperature for sampling
|
||||
"text_split_method": "cut0", # str. text split method, see text_segmentation_method.py for details.
|
||||
"text_split_method": "cut1", # str. text split method, see text_segmentation_method.py for details.
|
||||
"batch_size": 1, # int. batch size for inference
|
||||
"batch_threshold": 0.75, # float. threshold for batch splitting.
|
||||
"split_bucket: True, # bool. whether to split the batch into multiple buckets.
|
||||
"return_fragment": False, # bool. step by step return the audio fragment.
|
||||
"split_bucket": True, # bool. whether to split the batch into multiple buckets.
|
||||
"speed_factor":1.0, # float. control the speed of the synthesized audio.
|
||||
"fragment_interval":0.3, # float. to control the interval of the audio fragment.
|
||||
"seed": -1, # int. random seed for reproducibility.
|
||||
"parallel_infer": True, # bool. whether to use parallel inference.
|
||||
"repetition_penalty": 1.35 # float. repetition penalty for T2S model.
|
||||
"repetition_penalty": 1.35, # float. repetition penalty for T2S model.
|
||||
"sample_steps": 32, # int. number of sampling steps for VITS model V3.
|
||||
"super_sampling": False, # bool. whether to use super-sampling for audio when using VITS model V3.
|
||||
"super_sampling": False, # bool. whether to use super-sampling for audio when using VITS model V3.
|
||||
"return_fragment": False, # bool. step by step return the audio fragment. (Best Quality, Slowest response speed. old version of streaming mode)
|
||||
"streaming_mode": False, # bool. return audio chunk by chunk. (Medium quality, Slow response speed)
|
||||
"overlap_length": 2, # int. overlap length of semantic tokens for streaming mode.
|
||||
"min_chunk_length": 16, # int. The minimum chunk length of semantic tokens for streaming mode. (affects audio chunk size)
|
||||
"fixed_length_chunk": False, # bool. When turned on, it can achieve faster streaming response, but with lower quality. (lower quality, faster response speed)
|
||||
}
|
||||
returns:
|
||||
Tuple[int, np.ndarray]: sampling rate and audio data.
|
||||
@@ -1007,10 +1038,10 @@ class TTS:
|
||||
aux_ref_audio_paths: list = inputs.get("aux_ref_audio_paths", [])
|
||||
prompt_text: str = inputs.get("prompt_text", "")
|
||||
prompt_lang: str = inputs.get("prompt_lang", "")
|
||||
top_k: int = inputs.get("top_k", 5)
|
||||
top_k: int = inputs.get("top_k", 15)
|
||||
top_p: float = inputs.get("top_p", 1)
|
||||
temperature: float = inputs.get("temperature", 1)
|
||||
text_split_method: str = inputs.get("text_split_method", "cut0")
|
||||
text_split_method: str = inputs.get("text_split_method", "cut1")
|
||||
batch_size = inputs.get("batch_size", 1)
|
||||
batch_threshold = inputs.get("batch_threshold", 0.75)
|
||||
speed_factor = inputs.get("speed_factor", 1.0)
|
||||
@@ -1024,19 +1055,43 @@ class TTS:
|
||||
repetition_penalty = inputs.get("repetition_penalty", 1.35)
|
||||
sample_steps = inputs.get("sample_steps", 32)
|
||||
super_sampling = inputs.get("super_sampling", False)
|
||||
streaming_mode = inputs.get("streaming_mode", False)
|
||||
overlap_length = inputs.get("overlap_length", 2)
|
||||
min_chunk_length = inputs.get("min_chunk_length", 16)
|
||||
fixed_length_chunk = inputs.get("fixed_length_chunk", False)
|
||||
chunk_split_thershold = 0.0 # 该值代表语义token与mute token的余弦相似度阈值,若大于该阈值,则视为可切分点。
|
||||
|
||||
if parallel_infer:
|
||||
if parallel_infer and not streaming_mode:
|
||||
print(i18n("并行推理模式已开启"))
|
||||
self.t2s_model.model.infer_panel = self.t2s_model.model.infer_panel_batch_infer
|
||||
elif not parallel_infer and streaming_mode and not self.configs.use_vocoder:
|
||||
print(i18n("流式推理模式已开启"))
|
||||
self.t2s_model.model.infer_panel = self.t2s_model.model.infer_panel_naive
|
||||
elif streaming_mode and self.configs.use_vocoder:
|
||||
print(i18n("SoVits V3/4模型不支持流式推理模式,已自动回退到分段返回模式"))
|
||||
streaming_mode = False
|
||||
return_fragment = True
|
||||
if parallel_infer:
|
||||
self.t2s_model.model.infer_panel = self.t2s_model.model.infer_panel_batch_infer
|
||||
else:
|
||||
self.t2s_model.model.infer_panel = self.t2s_model.model.infer_panel_naive_batched
|
||||
# self.t2s_model.model.infer_panel = self.t2s_model.model.infer_panel_naive
|
||||
elif parallel_infer and streaming_mode:
|
||||
print(i18n("不支持同时开启并行推理和流式推理模式,已自动关闭并行推理模式"))
|
||||
parallel_infer = False
|
||||
self.t2s_model.model.infer_panel = self.t2s_model.model.infer_panel_naive
|
||||
else:
|
||||
print(i18n("并行推理模式已关闭"))
|
||||
print(i18n("朴素推理模式已开启"))
|
||||
self.t2s_model.model.infer_panel = self.t2s_model.model.infer_panel_naive_batched
|
||||
|
||||
if return_fragment:
|
||||
print(i18n("分段返回模式已开启"))
|
||||
if split_bucket:
|
||||
split_bucket = False
|
||||
print(i18n("分段返回模式不支持分桶处理,已自动关闭分桶处理"))
|
||||
if return_fragment and streaming_mode:
|
||||
print(i18n("流式推理模式不支持分段返回,已自动关闭分段返回"))
|
||||
return_fragment = False
|
||||
|
||||
if (return_fragment or streaming_mode) and split_bucket:
|
||||
print(i18n("分段返回模式/流式推理模式不支持分桶处理,已自动关闭分桶处理"))
|
||||
split_bucket = False
|
||||
|
||||
|
||||
if split_bucket and speed_factor == 1.0 and not (self.configs.use_vocoder and parallel_infer):
|
||||
print(i18n("分桶处理模式已开启"))
|
||||
@@ -1049,9 +1104,9 @@ class TTS:
|
||||
else:
|
||||
print(i18n("分桶处理模式已关闭"))
|
||||
|
||||
if fragment_interval < 0.01:
|
||||
fragment_interval = 0.01
|
||||
print(i18n("分段间隔过小,已自动设置为0.01"))
|
||||
# if fragment_interval < 0.01:
|
||||
# fragment_interval = 0.01
|
||||
# print(i18n("分段间隔过小,已自动设置为0.01"))
|
||||
|
||||
no_prompt_text = False
|
||||
if prompt_text in [None, ""]:
|
||||
@@ -1073,7 +1128,10 @@ class TTS:
|
||||
|
||||
###### setting reference audio and prompt text preprocessing ########
|
||||
t0 = time.perf_counter()
|
||||
if (ref_audio_path is not None) and (ref_audio_path != self.prompt_cache["ref_audio_path"]):
|
||||
if (ref_audio_path is not None) and (
|
||||
ref_audio_path != self.prompt_cache["ref_audio_path"]
|
||||
or (self.is_v2pro and self.prompt_cache["refer_spec"][0][1] is None)
|
||||
):
|
||||
if not os.path.exists(ref_audio_path):
|
||||
raise ValueError(f"{ref_audio_path} not exists")
|
||||
self.set_ref_audio(ref_audio_path)
|
||||
@@ -1109,7 +1167,7 @@ class TTS:
|
||||
###### text preprocessing ########
|
||||
t1 = time.perf_counter()
|
||||
data: list = None
|
||||
if not return_fragment:
|
||||
if not (return_fragment or streaming_mode):
|
||||
data = self.text_preprocessor.preprocess(text, text_lang, text_split_method, self.configs.version)
|
||||
if len(data) == 0:
|
||||
yield 16000, np.zeros(int(16000), dtype=np.int16)
|
||||
@@ -1169,10 +1227,11 @@ class TTS:
|
||||
t_34 = 0.0
|
||||
t_45 = 0.0
|
||||
audio = []
|
||||
is_first_package = True
|
||||
output_sr = self.configs.sampling_rate if not self.configs.use_vocoder else self.vocoder_configs["sr"]
|
||||
for item in data:
|
||||
t3 = time.perf_counter()
|
||||
if return_fragment:
|
||||
if return_fragment or streaming_mode:
|
||||
item = make_batch(item)
|
||||
if item is None:
|
||||
continue
|
||||
@@ -1194,99 +1253,228 @@ class TTS:
|
||||
self.prompt_cache["prompt_semantic"].expand(len(all_phoneme_ids), -1).to(self.configs.device)
|
||||
)
|
||||
|
||||
print(f"############ {i18n('预测语义Token')} ############")
|
||||
pred_semantic_list, idx_list = self.t2s_model.model.infer_panel(
|
||||
all_phoneme_ids,
|
||||
all_phoneme_lens,
|
||||
prompt,
|
||||
all_bert_features,
|
||||
# prompt_phone_len=ph_offset,
|
||||
top_k=top_k,
|
||||
top_p=top_p,
|
||||
temperature=temperature,
|
||||
early_stop_num=self.configs.hz * self.configs.max_sec,
|
||||
max_len=max_len,
|
||||
repetition_penalty=repetition_penalty,
|
||||
)
|
||||
t4 = time.perf_counter()
|
||||
t_34 += t4 - t3
|
||||
|
||||
refer_audio_spec = []
|
||||
if self.is_v2pro:sv_emb=[]
|
||||
for spec,audio_tensor in self.prompt_cache["refer_spec"]:
|
||||
spec=spec.to(dtype=self.precision, device=self.configs.device)
|
||||
|
||||
sv_emb = [] if self.is_v2pro else None
|
||||
for spec, audio_tensor in self.prompt_cache["refer_spec"]:
|
||||
spec = spec.to(dtype=self.precision, device=self.configs.device)
|
||||
refer_audio_spec.append(spec)
|
||||
if self.is_v2pro:
|
||||
sv_emb.append(self.sv_model.compute_embedding3(audio_tensor))
|
||||
|
||||
batch_audio_fragment = []
|
||||
if not streaming_mode:
|
||||
print(f"############ {i18n('预测语义Token')} ############")
|
||||
pred_semantic_list, idx_list = self.t2s_model.model.infer_panel(
|
||||
all_phoneme_ids,
|
||||
all_phoneme_lens,
|
||||
prompt,
|
||||
all_bert_features,
|
||||
# prompt_phone_len=ph_offset,
|
||||
top_k=top_k,
|
||||
top_p=top_p,
|
||||
temperature=temperature,
|
||||
early_stop_num=self.configs.hz * self.configs.max_sec,
|
||||
max_len=max_len,
|
||||
repetition_penalty=repetition_penalty,
|
||||
)
|
||||
t4 = time.perf_counter()
|
||||
t_34 += t4 - t3
|
||||
|
||||
# ## vits并行推理 method 1
|
||||
# pred_semantic_list = [item[-idx:] for item, idx in zip(pred_semantic_list, idx_list)]
|
||||
# pred_semantic_len = torch.LongTensor([item.shape[0] for item in pred_semantic_list]).to(self.configs.device)
|
||||
# pred_semantic = self.batch_sequences(pred_semantic_list, axis=0, pad_value=0).unsqueeze(0)
|
||||
# max_len = 0
|
||||
# for i in range(0, len(batch_phones)):
|
||||
# max_len = max(max_len, batch_phones[i].shape[-1])
|
||||
# batch_phones = self.batch_sequences(batch_phones, axis=0, pad_value=0, max_length=max_len)
|
||||
# batch_phones = batch_phones.to(self.configs.device)
|
||||
# batch_audio_fragment = (self.vits_model.batched_decode(
|
||||
# pred_semantic, pred_semantic_len, batch_phones, batch_phones_len,refer_audio_spec
|
||||
# ))
|
||||
print(f"############ {i18n('合成音频')} ############")
|
||||
if not self.configs.use_vocoder:
|
||||
if speed_factor == 1.0:
|
||||
print(f"{i18n('并行合成中')}...")
|
||||
# ## vits并行推理 method 2
|
||||
pred_semantic_list = [item[-idx:] for item, idx in zip(pred_semantic_list, idx_list)]
|
||||
upsample_rate = math.prod(self.vits_model.upsample_rates)
|
||||
audio_frag_idx = [
|
||||
pred_semantic_list[i].shape[0] * 2 * upsample_rate
|
||||
for i in range(0, len(pred_semantic_list))
|
||||
]
|
||||
audio_frag_end_idx = [sum(audio_frag_idx[: i + 1]) for i in range(0, len(audio_frag_idx))]
|
||||
all_pred_semantic = (
|
||||
torch.cat(pred_semantic_list).unsqueeze(0).unsqueeze(0).to(self.configs.device)
|
||||
)
|
||||
_batch_phones = torch.cat(batch_phones).unsqueeze(0).to(self.configs.device)
|
||||
if self.is_v2pro!=True:
|
||||
_batch_audio_fragment = self.vits_model.decode(all_pred_semantic, _batch_phones, refer_audio_spec, speed=speed_factor).detach()[0, 0, :]
|
||||
else:
|
||||
_batch_audio_fragment = self.vits_model.decode(all_pred_semantic, _batch_phones, refer_audio_spec, speed=speed_factor,sv_emb=sv_emb).detach()[0, 0, :]
|
||||
audio_frag_end_idx.insert(0, 0)
|
||||
batch_audio_fragment = [
|
||||
_batch_audio_fragment[audio_frag_end_idx[i - 1] : audio_frag_end_idx[i]]
|
||||
for i in range(1, len(audio_frag_end_idx))
|
||||
]
|
||||
else:
|
||||
# ## vits串行推理
|
||||
for i, idx in enumerate(tqdm(idx_list)):
|
||||
phones = batch_phones[i].unsqueeze(0).to(self.configs.device)
|
||||
_pred_semantic = (
|
||||
pred_semantic_list[i][-idx:].unsqueeze(0).unsqueeze(0)
|
||||
) # .unsqueeze(0)#mq要多unsqueeze一次
|
||||
if self.is_v2pro != True:
|
||||
audio_fragment = self.vits_model.decode(_pred_semantic, phones, refer_audio_spec, speed=speed_factor).detach()[0, 0, :]
|
||||
else:
|
||||
audio_fragment = self.vits_model.decode(_pred_semantic, phones, refer_audio_spec, speed=speed_factor,sv_emb=sv_emb).detach()[0, 0, :]
|
||||
batch_audio_fragment.append(audio_fragment) ###试试重建不带上prompt部分
|
||||
else:
|
||||
if parallel_infer:
|
||||
print(f"{i18n('并行合成中')}...")
|
||||
audio_fragments = self.using_vocoder_synthesis_batched_infer(
|
||||
idx_list, pred_semantic_list, batch_phones, speed=speed_factor, sample_steps=sample_steps
|
||||
)
|
||||
batch_audio_fragment.extend(audio_fragments)
|
||||
else:
|
||||
for i, idx in enumerate(tqdm(idx_list)):
|
||||
phones = batch_phones[i].unsqueeze(0).to(self.configs.device)
|
||||
_pred_semantic = (
|
||||
pred_semantic_list[i][-idx:].unsqueeze(0).unsqueeze(0)
|
||||
) # .unsqueeze(0)#mq要多unsqueeze一次
|
||||
audio_fragment = self.using_vocoder_synthesis(
|
||||
_pred_semantic, phones, speed=speed_factor, sample_steps=sample_steps
|
||||
|
||||
batch_audio_fragment = []
|
||||
|
||||
# ## vits并行推理 method 1
|
||||
# pred_semantic_list = [item[-idx:] for item, idx in zip(pred_semantic_list, idx_list)]
|
||||
# pred_semantic_len = torch.LongTensor([item.shape[0] for item in pred_semantic_list]).to(self.configs.device)
|
||||
# pred_semantic = self.batch_sequences(pred_semantic_list, axis=0, pad_value=0).unsqueeze(0)
|
||||
# max_len = 0
|
||||
# for i in range(0, len(batch_phones)):
|
||||
# max_len = max(max_len, batch_phones[i].shape[-1])
|
||||
# batch_phones = self.batch_sequences(batch_phones, axis=0, pad_value=0, max_length=max_len)
|
||||
# batch_phones = batch_phones.to(self.configs.device)
|
||||
# batch_audio_fragment = (self.vits_model.batched_decode(
|
||||
# pred_semantic, pred_semantic_len, batch_phones, batch_phones_len,refer_audio_spec
|
||||
# ))
|
||||
print(f"############ {i18n('合成音频')} ############")
|
||||
if not self.configs.use_vocoder:
|
||||
if speed_factor == 1.0:
|
||||
print(f"{i18n('并行合成中')}...")
|
||||
# ## vits并行推理 method 2
|
||||
pred_semantic_list = [item[-idx:] for item, idx in zip(pred_semantic_list, idx_list)]
|
||||
upsample_rate = math.prod(self.vits_model.upsample_rates)
|
||||
audio_frag_idx = [
|
||||
pred_semantic_list[i].shape[0] * 2 * upsample_rate
|
||||
for i in range(0, len(pred_semantic_list))
|
||||
]
|
||||
audio_frag_end_idx = [sum(audio_frag_idx[: i + 1]) for i in range(0, len(audio_frag_idx))]
|
||||
all_pred_semantic = (
|
||||
torch.cat(pred_semantic_list).unsqueeze(0).unsqueeze(0).to(self.configs.device)
|
||||
)
|
||||
batch_audio_fragment.append(audio_fragment)
|
||||
_batch_phones = torch.cat(batch_phones).unsqueeze(0).to(self.configs.device)
|
||||
|
||||
_batch_audio_fragment = self.vits_model.decode(
|
||||
all_pred_semantic, _batch_phones, refer_audio_spec, speed=speed_factor, sv_emb=sv_emb
|
||||
).detach()[0, 0, :]
|
||||
|
||||
audio_frag_end_idx.insert(0, 0)
|
||||
batch_audio_fragment = [
|
||||
_batch_audio_fragment[audio_frag_end_idx[i - 1] : audio_frag_end_idx[i]]
|
||||
for i in range(1, len(audio_frag_end_idx))
|
||||
]
|
||||
else:
|
||||
# ## vits串行推理
|
||||
for i, idx in enumerate(tqdm(idx_list)):
|
||||
phones = batch_phones[i].unsqueeze(0).to(self.configs.device)
|
||||
_pred_semantic = (
|
||||
pred_semantic_list[i][-idx:].unsqueeze(0).unsqueeze(0)
|
||||
) # .unsqueeze(0)#mq要多unsqueeze一次
|
||||
audio_fragment = self.vits_model.decode(
|
||||
_pred_semantic, phones, refer_audio_spec, speed=speed_factor, sv_emb=sv_emb
|
||||
).detach()[0, 0, :]
|
||||
batch_audio_fragment.append(audio_fragment) ###试试重建不带上prompt部分
|
||||
else:
|
||||
if parallel_infer:
|
||||
print(f"{i18n('并行合成中')}...")
|
||||
audio_fragments = self.using_vocoder_synthesis_batched_infer(
|
||||
idx_list, pred_semantic_list, batch_phones, speed=speed_factor, sample_steps=sample_steps
|
||||
)
|
||||
batch_audio_fragment.extend(audio_fragments)
|
||||
else:
|
||||
for i, idx in enumerate(tqdm(idx_list)):
|
||||
phones = batch_phones[i].unsqueeze(0).to(self.configs.device)
|
||||
_pred_semantic = (
|
||||
pred_semantic_list[i][-idx:].unsqueeze(0).unsqueeze(0)
|
||||
) # .unsqueeze(0)#mq要多unsqueeze一次
|
||||
audio_fragment = self.using_vocoder_synthesis(
|
||||
_pred_semantic, phones, speed=speed_factor, sample_steps=sample_steps
|
||||
)
|
||||
batch_audio_fragment.append(audio_fragment)
|
||||
|
||||
else:
|
||||
# refer_audio_spec: torch.Tensor = [
|
||||
# item.to(dtype=self.precision, device=self.configs.device)
|
||||
# for item in self.prompt_cache["refer_spec"]
|
||||
# ]
|
||||
semantic_token_generator =self.t2s_model.model.infer_panel(
|
||||
all_phoneme_ids[0].unsqueeze(0),
|
||||
all_phoneme_lens,
|
||||
prompt,
|
||||
all_bert_features[0].unsqueeze(0),
|
||||
top_k=top_k,
|
||||
top_p=top_p,
|
||||
temperature=temperature,
|
||||
early_stop_num=self.configs.hz * self.configs.max_sec,
|
||||
max_len=max_len,
|
||||
repetition_penalty=repetition_penalty,
|
||||
streaming_mode=True,
|
||||
chunk_length=min_chunk_length,
|
||||
mute_emb_sim_matrix=self.configs.mute_emb_sim_matrix if not fixed_length_chunk else None,
|
||||
chunk_split_thershold=chunk_split_thershold,
|
||||
)
|
||||
t4 = time.perf_counter()
|
||||
t_34 += t4 - t3
|
||||
phones = batch_phones[0].unsqueeze(0).to(self.configs.device)
|
||||
is_first_chunk = True
|
||||
|
||||
if not self.configs.use_vocoder:
|
||||
# if speed_factor == 1.0:
|
||||
# upsample_rate = math.prod(self.vits_model.upsample_rates)*(2 if self.vits_model.semantic_frame_rate == "25hz" else 1)
|
||||
# else:
|
||||
upsample_rate = math.prod(self.vits_model.upsample_rates)*((2 if self.vits_model.semantic_frame_rate == "25hz" else 1)/speed_factor)
|
||||
else:
|
||||
# if speed_factor == 1.0:
|
||||
# upsample_rate = self.vocoder_configs["upsample_rate"]*(3.875 if self.configs.version == "v3" else 4)
|
||||
# else:
|
||||
upsample_rate = self.vocoder_configs["upsample_rate"]*((3.875 if self.configs.version == "v3" else 4)/speed_factor)
|
||||
|
||||
last_audio_chunk = None
|
||||
# last_tokens = None
|
||||
last_latent = None
|
||||
previous_tokens = []
|
||||
overlap_len = overlap_length
|
||||
overlap_size = math.ceil(overlap_length*upsample_rate)
|
||||
for semantic_tokens, is_final in semantic_token_generator:
|
||||
if semantic_tokens is None and last_audio_chunk is not None:
|
||||
yield self.audio_postprocess(
|
||||
[[last_audio_chunk[-overlap_size:]]],
|
||||
output_sr,
|
||||
None,
|
||||
speed_factor,
|
||||
False,
|
||||
0.0,
|
||||
super_sampling if self.configs.use_vocoder and self.configs.version == "v3" else False,
|
||||
)
|
||||
break
|
||||
|
||||
_semantic_tokens = semantic_tokens
|
||||
print(f"semantic_tokens shape:{semantic_tokens.shape}")
|
||||
|
||||
previous_tokens.append(semantic_tokens)
|
||||
|
||||
_semantic_tokens = torch.cat(previous_tokens, dim=-1)
|
||||
|
||||
if not is_first_chunk and semantic_tokens.shape[-1] < 10:
|
||||
overlap_len = overlap_length+(10-semantic_tokens.shape[-1])
|
||||
else:
|
||||
overlap_len = overlap_length
|
||||
|
||||
|
||||
if not self.configs.use_vocoder:
|
||||
token_padding_length = 0
|
||||
# token_padding_length = int(phones.shape[-1]*2)-_semantic_tokens.shape[-1]
|
||||
# if token_padding_length>0:
|
||||
# _semantic_tokens = F.pad(_semantic_tokens, (0, token_padding_length), "constant", 486)
|
||||
# else:
|
||||
# token_padding_length = 0
|
||||
|
||||
audio_chunk, latent, latent_mask = self.vits_model.decode_streaming(
|
||||
_semantic_tokens.unsqueeze(0),
|
||||
phones, refer_audio_spec,
|
||||
speed=speed_factor,
|
||||
sv_emb=sv_emb,
|
||||
result_length=semantic_tokens.shape[-1]+overlap_len if not is_first_chunk else None,
|
||||
overlap_frames=last_latent[:,:,-overlap_len*(2 if self.vits_model.semantic_frame_rate == "25hz" else 1):] \
|
||||
if last_latent is not None else None,
|
||||
padding_length=token_padding_length
|
||||
)
|
||||
audio_chunk=audio_chunk.detach()[0, 0, :]
|
||||
else:
|
||||
raise RuntimeError(i18n("SoVits V3/4模型不支持流式推理模式"))
|
||||
|
||||
if overlap_len>overlap_length:
|
||||
audio_chunk=audio_chunk[-int((overlap_length+semantic_tokens.shape[-1])*upsample_rate):]
|
||||
|
||||
audio_chunk_ = audio_chunk
|
||||
if is_first_chunk and not is_final:
|
||||
is_first_chunk = False
|
||||
audio_chunk_ = audio_chunk_[:-overlap_size]
|
||||
elif is_first_chunk and is_final:
|
||||
is_first_chunk = False
|
||||
elif not is_first_chunk and not is_final:
|
||||
audio_chunk_ = self.sola_algorithm([last_audio_chunk, audio_chunk_], overlap_size)
|
||||
audio_chunk_ = (
|
||||
audio_chunk_[last_audio_chunk.shape[0]-overlap_size:-overlap_size] if not is_final \
|
||||
else audio_chunk_[last_audio_chunk.shape[0]-overlap_size:]
|
||||
)
|
||||
|
||||
last_latent = latent
|
||||
last_audio_chunk = audio_chunk
|
||||
yield self.audio_postprocess(
|
||||
[[audio_chunk_]],
|
||||
output_sr,
|
||||
None,
|
||||
speed_factor,
|
||||
False,
|
||||
0.0,
|
||||
super_sampling if self.configs.use_vocoder and self.configs.version == "v3" else False,
|
||||
)
|
||||
|
||||
if is_first_package:
|
||||
print(f"first_package_delay: {time.perf_counter()-t0:.3f}")
|
||||
is_first_package = False
|
||||
|
||||
|
||||
yield output_sr, np.zeros(int(output_sr*fragment_interval), dtype=np.int16)
|
||||
|
||||
t5 = time.perf_counter()
|
||||
t_45 += t5 - t4
|
||||
@@ -1301,17 +1489,18 @@ class TTS:
|
||||
fragment_interval,
|
||||
super_sampling if self.configs.use_vocoder and self.configs.version == "v3" else False,
|
||||
)
|
||||
elif streaming_mode:...
|
||||
else:
|
||||
audio.append(batch_audio_fragment)
|
||||
|
||||
if self.stop_flag:
|
||||
yield 16000, np.zeros(int(16000), dtype=np.int16)
|
||||
yield output_sr, np.zeros(int(output_sr), dtype=np.int16)
|
||||
return
|
||||
|
||||
if not return_fragment:
|
||||
if not (return_fragment or streaming_mode):
|
||||
print("%.3f\t%.3f\t%.3f\t%.3f" % (t1 - t0, t2 - t1, t_34, t_45))
|
||||
if len(audio) == 0:
|
||||
yield 16000, np.zeros(int(16000), dtype=np.int16)
|
||||
yield output_sr, np.zeros(int(output_sr), dtype=np.int16)
|
||||
return
|
||||
yield self.audio_postprocess(
|
||||
audio,
|
||||
@@ -1358,16 +1547,17 @@ class TTS:
|
||||
fragment_interval: float = 0.3,
|
||||
super_sampling: bool = False,
|
||||
) -> Tuple[int, np.ndarray]:
|
||||
zero_wav = torch.zeros(
|
||||
int(self.configs.sampling_rate * fragment_interval), dtype=self.precision, device=self.configs.device
|
||||
)
|
||||
if fragment_interval>0:
|
||||
zero_wav = torch.zeros(
|
||||
int(self.configs.sampling_rate * fragment_interval), dtype=self.precision, device=self.configs.device
|
||||
)
|
||||
|
||||
for i, batch in enumerate(audio):
|
||||
for j, audio_fragment in enumerate(batch):
|
||||
max_audio = torch.abs(audio_fragment).max() # 简单防止16bit爆音
|
||||
if max_audio > 1:
|
||||
audio_fragment /= max_audio
|
||||
audio_fragment: torch.Tensor = torch.cat([audio_fragment, zero_wav], dim=0)
|
||||
audio_fragment: torch.Tensor = torch.cat([audio_fragment, zero_wav], dim=0) if fragment_interval>0 else audio_fragment
|
||||
audio[i][j] = audio_fragment
|
||||
|
||||
if split_bucket:
|
||||
@@ -1387,12 +1577,16 @@ class TTS:
|
||||
max_audio = np.abs(audio).max()
|
||||
if max_audio > 1:
|
||||
audio /= max_audio
|
||||
audio = (audio * 32768).astype(np.int16)
|
||||
else:
|
||||
audio = audio.cpu().numpy()
|
||||
audio = (audio * 32768).astype(np.int16)
|
||||
t2 = time.perf_counter()
|
||||
print(f"超采样用时:{t2 - t1:.3f}s")
|
||||
else:
|
||||
audio = audio.cpu().numpy()
|
||||
audio = (audio * 32768).astype(np.int16)
|
||||
|
||||
audio = (audio * 32768).astype(np.int16)
|
||||
|
||||
# try:
|
||||
# if speed_factor != 1.0:
|
||||
@@ -1410,7 +1604,7 @@ class TTS:
|
||||
raw_entry = self.prompt_cache["refer_spec"][0]
|
||||
if isinstance(raw_entry, tuple):
|
||||
raw_entry = raw_entry[0]
|
||||
refer_audio_spec = raw_entry.to(dtype=self.precision,device=self.configs.device)
|
||||
refer_audio_spec = raw_entry.to(dtype=self.precision, device=self.configs.device)
|
||||
|
||||
fea_ref, ge = self.vits_model.decode_encp(prompt_semantic_tokens, prompt_phones, refer_audio_spec)
|
||||
ref_audio: torch.Tensor = self.prompt_cache["raw_audio"]
|
||||
@@ -1480,7 +1674,7 @@ class TTS:
|
||||
raw_entry = self.prompt_cache["refer_spec"][0]
|
||||
if isinstance(raw_entry, tuple):
|
||||
raw_entry = raw_entry[0]
|
||||
refer_audio_spec = raw_entry.to(dtype=self.precision,device=self.configs.device)
|
||||
refer_audio_spec = raw_entry.to(dtype=self.precision, device=self.configs.device)
|
||||
|
||||
fea_ref, ge = self.vits_model.decode_encp(prompt_semantic_tokens, prompt_phones, refer_audio_spec)
|
||||
ref_audio: torch.Tensor = self.prompt_cache["raw_audio"]
|
||||
@@ -1572,7 +1766,10 @@ class TTS:
|
||||
pos += chunk_len * upsample_rate
|
||||
|
||||
audio = self.sola_algorithm(audio_fragments, overlapped_len * upsample_rate)
|
||||
audio = audio[overlapped_len * upsample_rate : -padding_len * upsample_rate]
|
||||
if padding_len > 0:
|
||||
audio = audio[overlapped_len * upsample_rate : -padding_len * upsample_rate]
|
||||
else:
|
||||
audio = audio[overlapped_len * upsample_rate :]
|
||||
|
||||
audio_fragments = []
|
||||
for feat_len in feat_lens:
|
||||
@@ -1586,24 +1783,43 @@ class TTS:
|
||||
self,
|
||||
audio_fragments: List[torch.Tensor],
|
||||
overlap_len: int,
|
||||
search_len:int= 320
|
||||
):
|
||||
# overlap_len-=search_len
|
||||
|
||||
dtype = audio_fragments[0].dtype
|
||||
|
||||
for i in range(len(audio_fragments) - 1):
|
||||
f1 = audio_fragments[i]
|
||||
f2 = audio_fragments[i + 1]
|
||||
f1 = audio_fragments[i].float()
|
||||
f2 = audio_fragments[i + 1].float()
|
||||
w1 = f1[-overlap_len:]
|
||||
w2 = f2[:overlap_len]
|
||||
assert w1.shape == w2.shape
|
||||
corr = F.conv1d(w1.view(1, 1, -1), w2.view(1, 1, -1), padding=w2.shape[-1] // 2).view(-1)[:-1]
|
||||
idx = corr.argmax()
|
||||
f1_ = f1[: -(overlap_len - idx)]
|
||||
w2 = f2[:overlap_len+search_len]
|
||||
# w2 = w2[-w2.shape[-1]//2:]
|
||||
# assert w1.shape == w2.shape
|
||||
corr_norm = F.conv1d(w2.view(1, 1, -1), w1.view(1, 1, -1)).view(-1)
|
||||
|
||||
corr_den = F.conv1d(w2.view(1, 1, -1)**2, torch.ones_like(w1).view(1, 1, -1)).view(-1)+ 1e-8
|
||||
idx = (corr_norm/corr_den.sqrt()).argmax()
|
||||
|
||||
print(f"seg_idx: {idx}")
|
||||
|
||||
# idx = corr.argmax()
|
||||
f1_ = f1[: -overlap_len]
|
||||
audio_fragments[i] = f1_
|
||||
|
||||
f2_ = f2[idx:]
|
||||
window = torch.hann_window((overlap_len - idx) * 2, device=f1.device, dtype=f1.dtype)
|
||||
f2_[: (overlap_len - idx)] = (
|
||||
window[: (overlap_len - idx)] * f2_[: (overlap_len - idx)]
|
||||
+ window[(overlap_len - idx) :] * f1[-(overlap_len - idx) :]
|
||||
window = torch.hann_window((overlap_len) * 2, device=f1.device, dtype=f1.dtype)
|
||||
f2_[: overlap_len] = (
|
||||
window[: overlap_len] * f2_[: overlap_len]
|
||||
+ window[overlap_len :] * f1[-overlap_len :]
|
||||
)
|
||||
|
||||
# window = torch.sin(torch.arange((overlap_len - idx), device=f1.device) * np.pi / (overlap_len - idx))
|
||||
# f2_[: (overlap_len - idx)] = (
|
||||
# window * f2_[: (overlap_len - idx)]
|
||||
# + (1-window) * f1[-(overlap_len - idx) :]
|
||||
# )
|
||||
|
||||
audio_fragments[i + 1] = f2_
|
||||
|
||||
return torch.cat(audio_fragments, 0)
|
||||
return torch.cat(audio_fragments, 0).to(dtype)
|
||||
|
||||
@@ -121,65 +121,67 @@ class TextPreprocessor:
|
||||
|
||||
def get_phones_and_bert(self, text: str, language: str, version: str, final: bool = False):
|
||||
with self.bert_lock:
|
||||
if language in {"en", "all_zh", "all_ja", "all_ko", "all_yue"}:
|
||||
# language = language.replace("all_","")
|
||||
formattext = text
|
||||
while " " in formattext:
|
||||
formattext = formattext.replace(" ", " ")
|
||||
if language == "all_zh":
|
||||
if re.search(r"[A-Za-z]", formattext):
|
||||
formattext = re.sub(r"[a-z]", lambda x: x.group(0).upper(), formattext)
|
||||
formattext = chinese.mix_text_normalize(formattext)
|
||||
return self.get_phones_and_bert(formattext, "zh", version)
|
||||
text = re.sub(r' {2,}', ' ', text)
|
||||
textlist = []
|
||||
langlist = []
|
||||
if language == "all_zh":
|
||||
for tmp in LangSegmenter.getTexts(text,"zh"):
|
||||
langlist.append(tmp["lang"])
|
||||
textlist.append(tmp["text"])
|
||||
elif language == "all_yue":
|
||||
for tmp in LangSegmenter.getTexts(text,"zh"):
|
||||
if tmp["lang"] == "zh":
|
||||
tmp["lang"] = "yue"
|
||||
langlist.append(tmp["lang"])
|
||||
textlist.append(tmp["text"])
|
||||
elif language == "all_ja":
|
||||
for tmp in LangSegmenter.getTexts(text,"ja"):
|
||||
langlist.append(tmp["lang"])
|
||||
textlist.append(tmp["text"])
|
||||
elif language == "all_ko":
|
||||
for tmp in LangSegmenter.getTexts(text,"ko"):
|
||||
langlist.append(tmp["lang"])
|
||||
textlist.append(tmp["text"])
|
||||
elif language == "en":
|
||||
langlist.append("en")
|
||||
textlist.append(text)
|
||||
elif language == "auto":
|
||||
for tmp in LangSegmenter.getTexts(text):
|
||||
langlist.append(tmp["lang"])
|
||||
textlist.append(tmp["text"])
|
||||
elif language == "auto_yue":
|
||||
for tmp in LangSegmenter.getTexts(text):
|
||||
if tmp["lang"] == "zh":
|
||||
tmp["lang"] = "yue"
|
||||
langlist.append(tmp["lang"])
|
||||
textlist.append(tmp["text"])
|
||||
else:
|
||||
for tmp in LangSegmenter.getTexts(text):
|
||||
if langlist:
|
||||
if (tmp["lang"] == "en" and langlist[-1] == "en") or (tmp["lang"] != "en" and langlist[-1] != "en"):
|
||||
textlist[-1] += tmp["text"]
|
||||
continue
|
||||
if tmp["lang"] == "en":
|
||||
langlist.append(tmp["lang"])
|
||||
else:
|
||||
phones, word2ph, norm_text = self.clean_text_inf(formattext, language, version)
|
||||
bert = self.get_bert_feature(norm_text, word2ph).to(self.device)
|
||||
elif language == "all_yue" and re.search(r"[A-Za-z]", formattext):
|
||||
formattext = re.sub(r"[a-z]", lambda x: x.group(0).upper(), formattext)
|
||||
formattext = chinese.mix_text_normalize(formattext)
|
||||
return self.get_phones_and_bert(formattext, "yue", version)
|
||||
else:
|
||||
phones, word2ph, norm_text = self.clean_text_inf(formattext, language, version)
|
||||
bert = torch.zeros(
|
||||
(1024, len(phones)),
|
||||
dtype=torch.float32,
|
||||
).to(self.device)
|
||||
elif language in {"zh", "ja", "ko", "yue", "auto", "auto_yue"}:
|
||||
textlist = []
|
||||
langlist = []
|
||||
if language == "auto":
|
||||
for tmp in LangSegmenter.getTexts(text):
|
||||
langlist.append(tmp["lang"])
|
||||
textlist.append(tmp["text"])
|
||||
elif language == "auto_yue":
|
||||
for tmp in LangSegmenter.getTexts(text):
|
||||
if tmp["lang"] == "zh":
|
||||
tmp["lang"] = "yue"
|
||||
langlist.append(tmp["lang"])
|
||||
textlist.append(tmp["text"])
|
||||
else:
|
||||
for tmp in LangSegmenter.getTexts(text):
|
||||
if tmp["lang"] == "en":
|
||||
langlist.append(tmp["lang"])
|
||||
else:
|
||||
# 因无法区别中日韩文汉字,以用户输入为准
|
||||
langlist.append(language)
|
||||
textlist.append(tmp["text"])
|
||||
# print(textlist)
|
||||
# print(langlist)
|
||||
phones_list = []
|
||||
bert_list = []
|
||||
norm_text_list = []
|
||||
for i in range(len(textlist)):
|
||||
lang = langlist[i]
|
||||
phones, word2ph, norm_text = self.clean_text_inf(textlist[i], lang, version)
|
||||
bert = self.get_bert_inf(phones, word2ph, norm_text, lang)
|
||||
phones_list.append(phones)
|
||||
norm_text_list.append(norm_text)
|
||||
bert_list.append(bert)
|
||||
bert = torch.cat(bert_list, dim=1)
|
||||
phones = sum(phones_list, [])
|
||||
norm_text = "".join(norm_text_list)
|
||||
# 因无法区别中日韩文汉字,以用户输入为准
|
||||
langlist.append(language)
|
||||
textlist.append(tmp["text"])
|
||||
# print(textlist)
|
||||
# print(langlist)
|
||||
phones_list = []
|
||||
bert_list = []
|
||||
norm_text_list = []
|
||||
for i in range(len(textlist)):
|
||||
lang = langlist[i]
|
||||
phones, word2ph, norm_text = self.clean_text_inf(textlist[i], lang, version)
|
||||
bert = self.get_bert_inf(phones, word2ph, norm_text, lang)
|
||||
phones_list.append(phones)
|
||||
norm_text_list.append(norm_text)
|
||||
bert_list.append(bert)
|
||||
bert = torch.cat(bert_list, dim=1)
|
||||
phones = sum(phones_list, [])
|
||||
norm_text = "".join(norm_text_list)
|
||||
|
||||
if not final and len(phones) < 6:
|
||||
return self.get_phones_and_bert("." + text, language, version, final=True)
|
||||
@@ -234,4 +236,4 @@ class TextPreprocessor:
|
||||
punctuations = "".join(re.escape(p) for p in punctuation)
|
||||
pattern = f"([{punctuations}])([{punctuations}])+"
|
||||
result = re.sub(pattern, r"\1", text)
|
||||
return result
|
||||
return result
|
||||
@@ -92,7 +92,7 @@ def cut0(inp):
|
||||
if not set(inp).issubset(punctuation):
|
||||
return inp
|
||||
else:
|
||||
return "/n"
|
||||
return "\n"
|
||||
|
||||
|
||||
# 凑四句一切
|
||||
|
||||
@@ -22,6 +22,22 @@ v2:
|
||||
t2s_weights_path: GPT_SoVITS/pretrained_models/gsv-v2final-pretrained/s1bert25hz-5kh-longer-epoch=12-step=369668.ckpt
|
||||
version: v2
|
||||
vits_weights_path: GPT_SoVITS/pretrained_models/gsv-v2final-pretrained/s2G2333k.pth
|
||||
v2Pro:
|
||||
bert_base_path: GPT_SoVITS/pretrained_models/chinese-roberta-wwm-ext-large
|
||||
cnhuhbert_base_path: GPT_SoVITS/pretrained_models/chinese-hubert-base
|
||||
device: cpu
|
||||
is_half: false
|
||||
t2s_weights_path: GPT_SoVITS/pretrained_models/s1v3.ckpt
|
||||
version: v2Pro
|
||||
vits_weights_path: GPT_SoVITS/pretrained_models/v2Pro/s2Gv2Pro.pth
|
||||
v2ProPlus:
|
||||
bert_base_path: GPT_SoVITS/pretrained_models/chinese-roberta-wwm-ext-large
|
||||
cnhuhbert_base_path: GPT_SoVITS/pretrained_models/chinese-hubert-base
|
||||
device: cpu
|
||||
is_half: false
|
||||
t2s_weights_path: GPT_SoVITS/pretrained_models/s1v3.ckpt
|
||||
version: v2ProPlus
|
||||
vits_weights_path: GPT_SoVITS/pretrained_models/v2Pro/s2Gv2ProPlus.pth
|
||||
v3:
|
||||
bert_base_path: GPT_SoVITS/pretrained_models/chinese-roberta-wwm-ext-large
|
||||
cnhuhbert_base_path: GPT_SoVITS/pretrained_models/chinese-hubert-base
|
||||
|
||||
@@ -1,260 +0,0 @@
|
||||
# Copyright 3D-Speaker (https://github.com/alibaba-damo-academy/3D-Speaker). All Rights Reserved.
|
||||
# Licensed under the Apache License, Version 2.0 (http://www.apache.org/licenses/LICENSE-2.0)
|
||||
|
||||
"""
|
||||
Res2Net implementation is adapted from https://github.com/wenet-e2e/wespeaker.
|
||||
ERes2Net incorporates both local and global feature fusion techniques to improve the performance.
|
||||
The local feature fusion (LFF) fuses the features within one single residual block to extract the local signal.
|
||||
The global feature fusion (GFF) takes acoustic features of different scales as input to aggregate global signal.
|
||||
"""
|
||||
|
||||
|
||||
import torch
|
||||
import math
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
import pooling_layers as pooling_layers
|
||||
from fusion import AFF
|
||||
|
||||
class ReLU(nn.Hardtanh):
|
||||
|
||||
def __init__(self, inplace=False):
|
||||
super(ReLU, self).__init__(0, 20, inplace)
|
||||
|
||||
def __repr__(self):
|
||||
inplace_str = 'inplace' if self.inplace else ''
|
||||
return self.__class__.__name__ + ' (' \
|
||||
+ inplace_str + ')'
|
||||
|
||||
|
||||
class BasicBlockERes2Net(nn.Module):
|
||||
expansion = 2
|
||||
|
||||
def __init__(self, in_planes, planes, stride=1, baseWidth=32, scale=2):
|
||||
super(BasicBlockERes2Net, self).__init__()
|
||||
width = int(math.floor(planes*(baseWidth/64.0)))
|
||||
self.conv1 = nn.Conv2d(in_planes, width*scale, kernel_size=1, stride=stride, bias=False)
|
||||
self.bn1 = nn.BatchNorm2d(width*scale)
|
||||
self.nums = scale
|
||||
|
||||
convs=[]
|
||||
bns=[]
|
||||
for i in range(self.nums):
|
||||
convs.append(nn.Conv2d(width, width, kernel_size=3, padding=1, bias=False))
|
||||
bns.append(nn.BatchNorm2d(width))
|
||||
self.convs = nn.ModuleList(convs)
|
||||
self.bns = nn.ModuleList(bns)
|
||||
self.relu = ReLU(inplace=True)
|
||||
|
||||
self.conv3 = nn.Conv2d(width*scale, planes*self.expansion, kernel_size=1, bias=False)
|
||||
self.bn3 = nn.BatchNorm2d(planes*self.expansion)
|
||||
self.shortcut = nn.Sequential()
|
||||
if stride != 1 or in_planes != self.expansion * planes:
|
||||
self.shortcut = nn.Sequential(
|
||||
nn.Conv2d(in_planes, self.expansion * planes, kernel_size=1,
|
||||
stride=stride, bias=False),
|
||||
nn.BatchNorm2d(self.expansion * planes))
|
||||
self.stride = stride
|
||||
self.width = width
|
||||
self.scale = scale
|
||||
|
||||
def forward(self, x):
|
||||
residual = x
|
||||
|
||||
out = self.conv1(x)
|
||||
out = self.bn1(out)
|
||||
out = self.relu(out)
|
||||
spx = torch.split(out,self.width,1)
|
||||
for i in range(self.nums):
|
||||
if i==0:
|
||||
sp = spx[i]
|
||||
else:
|
||||
sp = sp + spx[i]
|
||||
sp = self.convs[i](sp)
|
||||
sp = self.relu(self.bns[i](sp))
|
||||
if i==0:
|
||||
out = sp
|
||||
else:
|
||||
out = torch.cat((out,sp),1)
|
||||
|
||||
out = self.conv3(out)
|
||||
out = self.bn3(out)
|
||||
|
||||
residual = self.shortcut(x)
|
||||
out += residual
|
||||
out = self.relu(out)
|
||||
|
||||
return out
|
||||
|
||||
class BasicBlockERes2Net_diff_AFF(nn.Module):
|
||||
expansion = 2
|
||||
|
||||
def __init__(self, in_planes, planes, stride=1, baseWidth=32, scale=2):
|
||||
super(BasicBlockERes2Net_diff_AFF, self).__init__()
|
||||
width = int(math.floor(planes*(baseWidth/64.0)))
|
||||
self.conv1 = nn.Conv2d(in_planes, width*scale, kernel_size=1, stride=stride, bias=False)
|
||||
self.bn1 = nn.BatchNorm2d(width*scale)
|
||||
self.nums = scale
|
||||
|
||||
convs=[]
|
||||
fuse_models=[]
|
||||
bns=[]
|
||||
for i in range(self.nums):
|
||||
convs.append(nn.Conv2d(width, width, kernel_size=3, padding=1, bias=False))
|
||||
bns.append(nn.BatchNorm2d(width))
|
||||
for j in range(self.nums - 1):
|
||||
fuse_models.append(AFF(channels=width))
|
||||
|
||||
self.convs = nn.ModuleList(convs)
|
||||
self.bns = nn.ModuleList(bns)
|
||||
self.fuse_models = nn.ModuleList(fuse_models)
|
||||
self.relu = ReLU(inplace=True)
|
||||
|
||||
self.conv3 = nn.Conv2d(width*scale, planes*self.expansion, kernel_size=1, bias=False)
|
||||
self.bn3 = nn.BatchNorm2d(planes*self.expansion)
|
||||
self.shortcut = nn.Sequential()
|
||||
if stride != 1 or in_planes != self.expansion * planes:
|
||||
self.shortcut = nn.Sequential(
|
||||
nn.Conv2d(in_planes, self.expansion * planes, kernel_size=1,
|
||||
stride=stride, bias=False),
|
||||
nn.BatchNorm2d(self.expansion * planes))
|
||||
self.stride = stride
|
||||
self.width = width
|
||||
self.scale = scale
|
||||
|
||||
def forward(self, x):
|
||||
residual = x
|
||||
|
||||
out = self.conv1(x)
|
||||
out = self.bn1(out)
|
||||
out = self.relu(out)
|
||||
spx = torch.split(out,self.width,1)
|
||||
for i in range(self.nums):
|
||||
if i==0:
|
||||
sp = spx[i]
|
||||
else:
|
||||
sp = self.fuse_models[i-1](sp, spx[i])
|
||||
|
||||
sp = self.convs[i](sp)
|
||||
sp = self.relu(self.bns[i](sp))
|
||||
if i==0:
|
||||
out = sp
|
||||
else:
|
||||
out = torch.cat((out,sp),1)
|
||||
|
||||
out = self.conv3(out)
|
||||
out = self.bn3(out)
|
||||
|
||||
residual = self.shortcut(x)
|
||||
out += residual
|
||||
out = self.relu(out)
|
||||
|
||||
return out
|
||||
|
||||
class ERes2Net(nn.Module):
|
||||
def __init__(self,
|
||||
block=BasicBlockERes2Net,
|
||||
block_fuse=BasicBlockERes2Net_diff_AFF,
|
||||
num_blocks=[3, 4, 6, 3],
|
||||
m_channels=32,
|
||||
feat_dim=80,
|
||||
embedding_size=192,
|
||||
pooling_func='TSTP',
|
||||
two_emb_layer=False):
|
||||
super(ERes2Net, self).__init__()
|
||||
self.in_planes = m_channels
|
||||
self.feat_dim = feat_dim
|
||||
self.embedding_size = embedding_size
|
||||
self.stats_dim = int(feat_dim / 8) * m_channels * 8
|
||||
self.two_emb_layer = two_emb_layer
|
||||
|
||||
self.conv1 = nn.Conv2d(1, m_channels, kernel_size=3, stride=1, padding=1, bias=False)
|
||||
self.bn1 = nn.BatchNorm2d(m_channels)
|
||||
self.layer1 = self._make_layer(block, m_channels, num_blocks[0], stride=1)
|
||||
self.layer2 = self._make_layer(block, m_channels * 2, num_blocks[1], stride=2)
|
||||
self.layer3 = self._make_layer(block_fuse, m_channels * 4, num_blocks[2], stride=2)
|
||||
self.layer4 = self._make_layer(block_fuse, m_channels * 8, num_blocks[3], stride=2)
|
||||
|
||||
# Downsampling module for each layer
|
||||
self.layer1_downsample = nn.Conv2d(m_channels * 2, m_channels * 4, kernel_size=3, stride=2, padding=1, bias=False)
|
||||
self.layer2_downsample = nn.Conv2d(m_channels * 4, m_channels * 8, kernel_size=3, padding=1, stride=2, bias=False)
|
||||
self.layer3_downsample = nn.Conv2d(m_channels * 8, m_channels * 16, kernel_size=3, padding=1, stride=2, bias=False)
|
||||
|
||||
# Bottom-up fusion module
|
||||
self.fuse_mode12 = AFF(channels=m_channels * 4)
|
||||
self.fuse_mode123 = AFF(channels=m_channels * 8)
|
||||
self.fuse_mode1234 = AFF(channels=m_channels * 16)
|
||||
|
||||
self.n_stats = 1 if pooling_func == 'TAP' or pooling_func == "TSDP" else 2
|
||||
self.pool = getattr(pooling_layers, pooling_func)(
|
||||
in_dim=self.stats_dim * block.expansion)
|
||||
self.seg_1 = nn.Linear(self.stats_dim * block.expansion * self.n_stats,
|
||||
embedding_size)
|
||||
if self.two_emb_layer:
|
||||
self.seg_bn_1 = nn.BatchNorm1d(embedding_size, affine=False)
|
||||
self.seg_2 = nn.Linear(embedding_size, embedding_size)
|
||||
else:
|
||||
self.seg_bn_1 = nn.Identity()
|
||||
self.seg_2 = nn.Identity()
|
||||
|
||||
def _make_layer(self, block, planes, num_blocks, stride):
|
||||
strides = [stride] + [1] * (num_blocks - 1)
|
||||
layers = []
|
||||
for stride in strides:
|
||||
layers.append(block(self.in_planes, planes, stride))
|
||||
self.in_planes = planes * block.expansion
|
||||
return nn.Sequential(*layers)
|
||||
|
||||
def forward(self, x):
|
||||
x = x.permute(0, 2, 1) # (B,T,F) => (B,F,T)
|
||||
x = x.unsqueeze_(1)
|
||||
out = F.relu(self.bn1(self.conv1(x)))
|
||||
out1 = self.layer1(out)
|
||||
out2 = self.layer2(out1)
|
||||
out1_downsample = self.layer1_downsample(out1)
|
||||
fuse_out12 = self.fuse_mode12(out2, out1_downsample)
|
||||
out3 = self.layer3(out2)
|
||||
fuse_out12_downsample = self.layer2_downsample(fuse_out12)
|
||||
fuse_out123 = self.fuse_mode123(out3, fuse_out12_downsample)
|
||||
out4 = self.layer4(out3)
|
||||
fuse_out123_downsample = self.layer3_downsample(fuse_out123)
|
||||
fuse_out1234 = self.fuse_mode1234(out4, fuse_out123_downsample)
|
||||
stats = self.pool(fuse_out1234)
|
||||
|
||||
embed_a = self.seg_1(stats)
|
||||
if self.two_emb_layer:
|
||||
out = F.relu(embed_a)
|
||||
out = self.seg_bn_1(out)
|
||||
embed_b = self.seg_2(out)
|
||||
return embed_b
|
||||
else:
|
||||
return embed_a
|
||||
|
||||
def forward3(self, x):
|
||||
x = x.permute(0, 2, 1) # (B,T,F) => (B,F,T)
|
||||
x = x.unsqueeze_(1)
|
||||
out = F.relu(self.bn1(self.conv1(x)))
|
||||
out1 = self.layer1(out)
|
||||
out2 = self.layer2(out1)
|
||||
out1_downsample = self.layer1_downsample(out1)
|
||||
fuse_out12 = self.fuse_mode12(out2, out1_downsample)
|
||||
out3 = self.layer3(out2)
|
||||
fuse_out12_downsample = self.layer2_downsample(fuse_out12)
|
||||
fuse_out123 = self.fuse_mode123(out3, fuse_out12_downsample)
|
||||
out4 = self.layer4(out3)
|
||||
fuse_out123_downsample = self.layer3_downsample(fuse_out123)
|
||||
fuse_out1234 = self.fuse_mode1234(out4, fuse_out123_downsample).flatten(start_dim=1,end_dim=2).mean(-1)
|
||||
return fuse_out1234
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
|
||||
x = torch.zeros(10, 300, 80)
|
||||
model = ERes2Net(feat_dim=80, embedding_size=192, pooling_func='TSTP')
|
||||
model.eval()
|
||||
out = model(x)
|
||||
print(out.shape) # torch.Size([10, 192])
|
||||
|
||||
num_params = sum(param.numel() for param in model.parameters())
|
||||
print("{} M".format(num_params / 1e6)) # 6.61M
|
||||
|
||||
@@ -1,14 +1,12 @@
|
||||
# Copyright 3D-Speaker (https://github.com/alibaba-damo-academy/3D-Speaker). All Rights Reserved.
|
||||
# Licensed under the Apache License, Version 2.0 (http://www.apache.org/licenses/LICENSE-2.0)
|
||||
|
||||
"""
|
||||
To further improve the short-duration feature extraction capability of ERes2Net, we expand the channel dimension
|
||||
within each stage. However, this modification also increases the number of model parameters and computational complexity.
|
||||
To alleviate this problem, we propose an improved ERes2NetV2 by pruning redundant structures, ultimately reducing
|
||||
both the model parameters and its computational cost.
|
||||
"""
|
||||
|
||||
|
||||
To further improve the short-duration feature extraction capability of ERes2Net, we expand the channel dimension
|
||||
within each stage. However, this modification also increases the number of model parameters and computational complexity.
|
||||
To alleviate this problem, we propose an improved ERes2NetV2 by pruning redundant structures, ultimately reducing
|
||||
both the model parameters and its computational cost.
|
||||
"""
|
||||
|
||||
import torch
|
||||
import math
|
||||
@@ -17,47 +15,42 @@ import torch.nn.functional as F
|
||||
import pooling_layers as pooling_layers
|
||||
from fusion import AFF
|
||||
|
||||
class ReLU(nn.Hardtanh):
|
||||
|
||||
class ReLU(nn.Hardtanh):
|
||||
def __init__(self, inplace=False):
|
||||
super(ReLU, self).__init__(0, 20, inplace)
|
||||
|
||||
def __repr__(self):
|
||||
inplace_str = 'inplace' if self.inplace else ''
|
||||
return self.__class__.__name__ + ' (' \
|
||||
+ inplace_str + ')'
|
||||
inplace_str = "inplace" if self.inplace else ""
|
||||
return self.__class__.__name__ + " (" + inplace_str + ")"
|
||||
|
||||
|
||||
class BasicBlockERes2NetV2(nn.Module):
|
||||
|
||||
def __init__(self, in_planes, planes, stride=1, baseWidth=26, scale=2, expansion=2):
|
||||
super(BasicBlockERes2NetV2, self).__init__()
|
||||
width = int(math.floor(planes*(baseWidth/64.0)))
|
||||
self.conv1 = nn.Conv2d(in_planes, width*scale, kernel_size=1, stride=stride, bias=False)
|
||||
self.bn1 = nn.BatchNorm2d(width*scale)
|
||||
width = int(math.floor(planes * (baseWidth / 64.0)))
|
||||
self.conv1 = nn.Conv2d(in_planes, width * scale, kernel_size=1, stride=stride, bias=False)
|
||||
self.bn1 = nn.BatchNorm2d(width * scale)
|
||||
self.nums = scale
|
||||
self.expansion = expansion
|
||||
|
||||
convs=[]
|
||||
bns=[]
|
||||
convs = []
|
||||
bns = []
|
||||
for i in range(self.nums):
|
||||
convs.append(nn.Conv2d(width, width, kernel_size=3, padding=1, bias=False))
|
||||
bns.append(nn.BatchNorm2d(width))
|
||||
convs.append(nn.Conv2d(width, width, kernel_size=3, padding=1, bias=False))
|
||||
bns.append(nn.BatchNorm2d(width))
|
||||
self.convs = nn.ModuleList(convs)
|
||||
self.bns = nn.ModuleList(bns)
|
||||
self.relu = ReLU(inplace=True)
|
||||
|
||||
self.conv3 = nn.Conv2d(width*scale, planes*self.expansion, kernel_size=1, bias=False)
|
||||
self.bn3 = nn.BatchNorm2d(planes*self.expansion)
|
||||
|
||||
self.conv3 = nn.Conv2d(width * scale, planes * self.expansion, kernel_size=1, bias=False)
|
||||
self.bn3 = nn.BatchNorm2d(planes * self.expansion)
|
||||
self.shortcut = nn.Sequential()
|
||||
if stride != 1 or in_planes != self.expansion * planes:
|
||||
self.shortcut = nn.Sequential(
|
||||
nn.Conv2d(in_planes,
|
||||
self.expansion * planes,
|
||||
kernel_size=1,
|
||||
stride=stride,
|
||||
bias=False),
|
||||
nn.BatchNorm2d(self.expansion * planes))
|
||||
nn.Conv2d(in_planes, self.expansion * planes, kernel_size=1, stride=stride, bias=False),
|
||||
nn.BatchNorm2d(self.expansion * planes),
|
||||
)
|
||||
self.stride = stride
|
||||
self.width = width
|
||||
self.scale = scale
|
||||
@@ -68,18 +61,18 @@ class BasicBlockERes2NetV2(nn.Module):
|
||||
out = self.conv1(x)
|
||||
out = self.bn1(out)
|
||||
out = self.relu(out)
|
||||
spx = torch.split(out,self.width,1)
|
||||
spx = torch.split(out, self.width, 1)
|
||||
for i in range(self.nums):
|
||||
if i==0:
|
||||
sp = spx[i]
|
||||
else:
|
||||
sp = sp + spx[i]
|
||||
sp = self.convs[i](sp)
|
||||
sp = self.relu(self.bns[i](sp))
|
||||
if i==0:
|
||||
out = sp
|
||||
else:
|
||||
out = torch.cat((out,sp),1)
|
||||
if i == 0:
|
||||
sp = spx[i]
|
||||
else:
|
||||
sp = sp + spx[i]
|
||||
sp = self.convs[i](sp)
|
||||
sp = self.relu(self.bns[i](sp))
|
||||
if i == 0:
|
||||
out = sp
|
||||
else:
|
||||
out = torch.cat((out, sp), 1)
|
||||
|
||||
out = self.conv3(out)
|
||||
out = self.bn3(out)
|
||||
@@ -90,22 +83,22 @@ class BasicBlockERes2NetV2(nn.Module):
|
||||
|
||||
return out
|
||||
|
||||
class BasicBlockERes2NetV2AFF(nn.Module):
|
||||
|
||||
class BasicBlockERes2NetV2AFF(nn.Module):
|
||||
def __init__(self, in_planes, planes, stride=1, baseWidth=26, scale=2, expansion=2):
|
||||
super(BasicBlockERes2NetV2AFF, self).__init__()
|
||||
width = int(math.floor(planes*(baseWidth/64.0)))
|
||||
self.conv1 = nn.Conv2d(in_planes, width*scale, kernel_size=1, stride=stride, bias=False)
|
||||
self.bn1 = nn.BatchNorm2d(width*scale)
|
||||
width = int(math.floor(planes * (baseWidth / 64.0)))
|
||||
self.conv1 = nn.Conv2d(in_planes, width * scale, kernel_size=1, stride=stride, bias=False)
|
||||
self.bn1 = nn.BatchNorm2d(width * scale)
|
||||
self.nums = scale
|
||||
self.expansion = expansion
|
||||
|
||||
convs=[]
|
||||
fuse_models=[]
|
||||
bns=[]
|
||||
convs = []
|
||||
fuse_models = []
|
||||
bns = []
|
||||
for i in range(self.nums):
|
||||
convs.append(nn.Conv2d(width, width, kernel_size=3, padding=1, bias=False))
|
||||
bns.append(nn.BatchNorm2d(width))
|
||||
convs.append(nn.Conv2d(width, width, kernel_size=3, padding=1, bias=False))
|
||||
bns.append(nn.BatchNorm2d(width))
|
||||
for j in range(self.nums - 1):
|
||||
fuse_models.append(AFF(channels=width, r=4))
|
||||
|
||||
@@ -113,18 +106,15 @@ class BasicBlockERes2NetV2AFF(nn.Module):
|
||||
self.bns = nn.ModuleList(bns)
|
||||
self.fuse_models = nn.ModuleList(fuse_models)
|
||||
self.relu = ReLU(inplace=True)
|
||||
|
||||
self.conv3 = nn.Conv2d(width*scale, planes*self.expansion, kernel_size=1, bias=False)
|
||||
self.bn3 = nn.BatchNorm2d(planes*self.expansion)
|
||||
|
||||
self.conv3 = nn.Conv2d(width * scale, planes * self.expansion, kernel_size=1, bias=False)
|
||||
self.bn3 = nn.BatchNorm2d(planes * self.expansion)
|
||||
self.shortcut = nn.Sequential()
|
||||
if stride != 1 or in_planes != self.expansion * planes:
|
||||
self.shortcut = nn.Sequential(
|
||||
nn.Conv2d(in_planes,
|
||||
self.expansion * planes,
|
||||
kernel_size=1,
|
||||
stride=stride,
|
||||
bias=False),
|
||||
nn.BatchNorm2d(self.expansion * planes))
|
||||
nn.Conv2d(in_planes, self.expansion * planes, kernel_size=1, stride=stride, bias=False),
|
||||
nn.BatchNorm2d(self.expansion * planes),
|
||||
)
|
||||
self.stride = stride
|
||||
self.width = width
|
||||
self.scale = scale
|
||||
@@ -135,19 +125,19 @@ class BasicBlockERes2NetV2AFF(nn.Module):
|
||||
out = self.conv1(x)
|
||||
out = self.bn1(out)
|
||||
out = self.relu(out)
|
||||
spx = torch.split(out,self.width,1)
|
||||
spx = torch.split(out, self.width, 1)
|
||||
for i in range(self.nums):
|
||||
if i==0:
|
||||
if i == 0:
|
||||
sp = spx[i]
|
||||
else:
|
||||
sp = self.fuse_models[i-1](sp, spx[i])
|
||||
|
||||
sp = self.fuse_models[i - 1](sp, spx[i])
|
||||
|
||||
sp = self.convs[i](sp)
|
||||
sp = self.relu(self.bns[i](sp))
|
||||
if i==0:
|
||||
if i == 0:
|
||||
out = sp
|
||||
else:
|
||||
out = torch.cat((out,sp),1)
|
||||
out = torch.cat((out, sp), 1)
|
||||
|
||||
out = self.conv3(out)
|
||||
out = self.bn3(out)
|
||||
@@ -158,19 +148,22 @@ class BasicBlockERes2NetV2AFF(nn.Module):
|
||||
|
||||
return out
|
||||
|
||||
|
||||
class ERes2NetV2(nn.Module):
|
||||
def __init__(self,
|
||||
block=BasicBlockERes2NetV2,
|
||||
block_fuse=BasicBlockERes2NetV2AFF,
|
||||
num_blocks=[3, 4, 6, 3],
|
||||
m_channels=64,
|
||||
feat_dim=80,
|
||||
embedding_size=192,
|
||||
baseWidth=26,
|
||||
scale=2,
|
||||
expansion=2,
|
||||
pooling_func='TSTP',
|
||||
two_emb_layer=False):
|
||||
def __init__(
|
||||
self,
|
||||
block=BasicBlockERes2NetV2,
|
||||
block_fuse=BasicBlockERes2NetV2AFF,
|
||||
num_blocks=[3, 4, 6, 3],
|
||||
m_channels=64,
|
||||
feat_dim=80,
|
||||
embedding_size=192,
|
||||
baseWidth=26,
|
||||
scale=2,
|
||||
expansion=2,
|
||||
pooling_func="TSTP",
|
||||
two_emb_layer=False,
|
||||
):
|
||||
super(ERes2NetV2, self).__init__()
|
||||
self.in_planes = m_channels
|
||||
self.feat_dim = feat_dim
|
||||
@@ -181,42 +174,29 @@ class ERes2NetV2(nn.Module):
|
||||
self.scale = scale
|
||||
self.expansion = expansion
|
||||
|
||||
self.conv1 = nn.Conv2d(1,
|
||||
m_channels,
|
||||
kernel_size=3,
|
||||
stride=1,
|
||||
padding=1,
|
||||
bias=False)
|
||||
self.conv1 = nn.Conv2d(1, m_channels, kernel_size=3, stride=1, padding=1, bias=False)
|
||||
self.bn1 = nn.BatchNorm2d(m_channels)
|
||||
self.layer1 = self._make_layer(block,
|
||||
m_channels,
|
||||
num_blocks[0],
|
||||
stride=1)
|
||||
self.layer2 = self._make_layer(block,
|
||||
m_channels * 2,
|
||||
num_blocks[1],
|
||||
stride=2)
|
||||
self.layer3 = self._make_layer(block_fuse,
|
||||
m_channels * 4,
|
||||
num_blocks[2],
|
||||
stride=2)
|
||||
self.layer4 = self._make_layer(block_fuse,
|
||||
m_channels * 8,
|
||||
num_blocks[3],
|
||||
stride=2)
|
||||
self.layer1 = self._make_layer(block, m_channels, num_blocks[0], stride=1)
|
||||
self.layer2 = self._make_layer(block, m_channels * 2, num_blocks[1], stride=2)
|
||||
self.layer3 = self._make_layer(block_fuse, m_channels * 4, num_blocks[2], stride=2)
|
||||
self.layer4 = self._make_layer(block_fuse, m_channels * 8, num_blocks[3], stride=2)
|
||||
|
||||
# Downsampling module
|
||||
self.layer3_ds = nn.Conv2d(m_channels * 4 * self.expansion, m_channels * 8 * self.expansion, kernel_size=3, \
|
||||
padding=1, stride=2, bias=False)
|
||||
self.layer3_ds = nn.Conv2d(
|
||||
m_channels * 4 * self.expansion,
|
||||
m_channels * 8 * self.expansion,
|
||||
kernel_size=3,
|
||||
padding=1,
|
||||
stride=2,
|
||||
bias=False,
|
||||
)
|
||||
|
||||
# Bottom-up fusion module
|
||||
self.fuse34 = AFF(channels=m_channels * 8 * self.expansion, r=4)
|
||||
|
||||
self.n_stats = 1 if pooling_func == 'TAP' or pooling_func == "TSDP" else 2
|
||||
self.pool = getattr(pooling_layers, pooling_func)(
|
||||
in_dim=self.stats_dim * self.expansion)
|
||||
self.seg_1 = nn.Linear(self.stats_dim * self.expansion * self.n_stats,
|
||||
embedding_size)
|
||||
self.n_stats = 1 if pooling_func == "TAP" or pooling_func == "TSDP" else 2
|
||||
self.pool = getattr(pooling_layers, pooling_func)(in_dim=self.stats_dim * self.expansion)
|
||||
self.seg_1 = nn.Linear(self.stats_dim * self.expansion * self.n_stats, embedding_size)
|
||||
if self.two_emb_layer:
|
||||
self.seg_bn_1 = nn.BatchNorm1d(embedding_size, affine=False)
|
||||
self.seg_2 = nn.Linear(embedding_size, embedding_size)
|
||||
@@ -228,7 +208,11 @@ class ERes2NetV2(nn.Module):
|
||||
strides = [stride] + [1] * (num_blocks - 1)
|
||||
layers = []
|
||||
for stride in strides:
|
||||
layers.append(block(self.in_planes, planes, stride, baseWidth=self.baseWidth, scale=self.scale, expansion=self.expansion))
|
||||
layers.append(
|
||||
block(
|
||||
self.in_planes, planes, stride, baseWidth=self.baseWidth, scale=self.scale, expansion=self.expansion
|
||||
)
|
||||
)
|
||||
self.in_planes = planes * self.expansion
|
||||
return nn.Sequential(*layers)
|
||||
|
||||
@@ -264,7 +248,7 @@ class ERes2NetV2(nn.Module):
|
||||
out3_ds = self.layer3_ds(out3)
|
||||
fuse_out34 = self.fuse34(out4, out3_ds)
|
||||
# print(111111111,fuse_out34.shape)#111111111 torch.Size([16, 2048, 10, 72])
|
||||
return fuse_out34.flatten(start_dim=1,end_dim=2).mean(-1)
|
||||
return fuse_out34.flatten(start_dim=1, end_dim=2).mean(-1)
|
||||
# stats = self.pool(fuse_out34)
|
||||
#
|
||||
# embed_a = self.seg_1(stats)
|
||||
@@ -275,18 +259,3 @@ class ERes2NetV2(nn.Module):
|
||||
# return embed_b
|
||||
# else:
|
||||
# return embed_a
|
||||
|
||||
if __name__ == '__main__':
|
||||
|
||||
x = torch.randn(1, 300, 80)
|
||||
model = ERes2NetV2(feat_dim=80, embedding_size=192, m_channels=64, baseWidth=26, scale=2, expansion=2)
|
||||
model.eval()
|
||||
y = model(x)
|
||||
print(y.size())
|
||||
macs, num_params = profile(model, inputs=(x, ))
|
||||
print("Params: {} M".format(num_params / 1e6)) # 17.86 M
|
||||
print("MACs: {} G".format(macs / 1e9)) # 12.69 G
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
@@ -1,286 +0,0 @@
|
||||
# Copyright 3D-Speaker (https://github.com/alibaba-damo-academy/3D-Speaker). All Rights Reserved.
|
||||
# Licensed under the Apache License, Version 2.0 (http://www.apache.org/licenses/LICENSE-2.0)
|
||||
|
||||
""" Res2Net implementation is adapted from https://github.com/wenet-e2e/wespeaker.
|
||||
ERes2Net incorporates both local and global feature fusion techniques to improve the performance.
|
||||
The local feature fusion (LFF) fuses the features within one single residual block to extract the local signal.
|
||||
The global feature fusion (GFF) takes acoustic features of different scales as input to aggregate global signal.
|
||||
ERes2Net-huge is an upgraded version of ERes2Net that uses a larger number of parameters to achieve better
|
||||
recognition performance. Parameters expansion, baseWidth, and scale can be modified to obtain optimal performance.
|
||||
"""
|
||||
import pdb
|
||||
|
||||
import torch
|
||||
import math
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
import pooling_layers as pooling_layers
|
||||
from fusion import AFF
|
||||
|
||||
class ReLU(nn.Hardtanh):
|
||||
|
||||
def __init__(self, inplace=False):
|
||||
super(ReLU, self).__init__(0, 20, inplace)
|
||||
|
||||
def __repr__(self):
|
||||
inplace_str = 'inplace' if self.inplace else ''
|
||||
return self.__class__.__name__ + ' (' \
|
||||
+ inplace_str + ')'
|
||||
|
||||
|
||||
class BasicBlockERes2Net(nn.Module):
|
||||
expansion = 4
|
||||
|
||||
def __init__(self, in_planes, planes, stride=1, baseWidth=24, scale=3):
|
||||
super(BasicBlockERes2Net, self).__init__()
|
||||
width = int(math.floor(planes*(baseWidth/64.0)))
|
||||
self.conv1 = nn.Conv2d(in_planes, width*scale, kernel_size=1, stride=stride, bias=False)
|
||||
self.bn1 = nn.BatchNorm2d(width*scale)
|
||||
self.nums = scale
|
||||
|
||||
convs=[]
|
||||
bns=[]
|
||||
for i in range(self.nums):
|
||||
convs.append(nn.Conv2d(width, width, kernel_size=3, padding=1, bias=False))
|
||||
bns.append(nn.BatchNorm2d(width))
|
||||
self.convs = nn.ModuleList(convs)
|
||||
self.bns = nn.ModuleList(bns)
|
||||
self.relu = ReLU(inplace=True)
|
||||
|
||||
self.conv3 = nn.Conv2d(width*scale, planes*self.expansion, kernel_size=1, bias=False)
|
||||
self.bn3 = nn.BatchNorm2d(planes*self.expansion)
|
||||
self.shortcut = nn.Sequential()
|
||||
if stride != 1 or in_planes != self.expansion * planes:
|
||||
self.shortcut = nn.Sequential(
|
||||
nn.Conv2d(in_planes, self.expansion * planes, kernel_size=1, stride=stride, bias=False),
|
||||
nn.BatchNorm2d(self.expansion * planes))
|
||||
self.stride = stride
|
||||
self.width = width
|
||||
self.scale = scale
|
||||
|
||||
def forward(self, x):
|
||||
residual = x
|
||||
|
||||
out = self.conv1(x)
|
||||
out = self.bn1(out)
|
||||
out = self.relu(out)
|
||||
spx = torch.split(out,self.width,1)
|
||||
for i in range(self.nums):
|
||||
if i==0:
|
||||
sp = spx[i]
|
||||
else:
|
||||
sp = sp + spx[i]
|
||||
sp = self.convs[i](sp)
|
||||
sp = self.relu(self.bns[i](sp))
|
||||
if i==0:
|
||||
out = sp
|
||||
else:
|
||||
out = torch.cat((out,sp),1)
|
||||
|
||||
out = self.conv3(out)
|
||||
out = self.bn3(out)
|
||||
|
||||
residual = self.shortcut(x)
|
||||
out += residual
|
||||
out = self.relu(out)
|
||||
|
||||
return out
|
||||
|
||||
class BasicBlockERes2Net_diff_AFF(nn.Module):
|
||||
expansion = 4
|
||||
|
||||
def __init__(self, in_planes, planes, stride=1, baseWidth=24, scale=3):
|
||||
super(BasicBlockERes2Net_diff_AFF, self).__init__()
|
||||
width = int(math.floor(planes*(baseWidth/64.0)))
|
||||
self.conv1 = nn.Conv2d(in_planes, width*scale, kernel_size=1, stride=stride, bias=False)
|
||||
self.bn1 = nn.BatchNorm2d(width*scale)
|
||||
self.nums = scale
|
||||
|
||||
convs=[]
|
||||
fuse_models=[]
|
||||
bns=[]
|
||||
for i in range(self.nums):
|
||||
convs.append(nn.Conv2d(width, width, kernel_size=3, padding=1, bias=False))
|
||||
bns.append(nn.BatchNorm2d(width))
|
||||
for j in range(self.nums - 1):
|
||||
fuse_models.append(AFF(channels=width))
|
||||
|
||||
self.convs = nn.ModuleList(convs)
|
||||
self.bns = nn.ModuleList(bns)
|
||||
self.fuse_models = nn.ModuleList(fuse_models)
|
||||
self.relu = ReLU(inplace=True)
|
||||
|
||||
self.conv3 = nn.Conv2d(width*scale, planes*self.expansion, kernel_size=1, bias=False)
|
||||
self.bn3 = nn.BatchNorm2d(planes*self.expansion)
|
||||
self.shortcut = nn.Sequential()
|
||||
if stride != 1 or in_planes != self.expansion * planes:
|
||||
self.shortcut = nn.Sequential(
|
||||
nn.Conv2d(in_planes, self.expansion * planes, kernel_size=1, stride=stride, bias=False),
|
||||
nn.BatchNorm2d(self.expansion * planes))
|
||||
self.stride = stride
|
||||
self.width = width
|
||||
self.scale = scale
|
||||
|
||||
def forward(self, x):
|
||||
residual = x
|
||||
|
||||
out = self.conv1(x)
|
||||
out = self.bn1(out)
|
||||
out = self.relu(out)
|
||||
spx = torch.split(out,self.width,1)
|
||||
for i in range(self.nums):
|
||||
if i==0:
|
||||
sp = spx[i]
|
||||
else:
|
||||
sp = self.fuse_models[i-1](sp, spx[i])
|
||||
|
||||
sp = self.convs[i](sp)
|
||||
sp = self.relu(self.bns[i](sp))
|
||||
if i==0:
|
||||
out = sp
|
||||
else:
|
||||
out = torch.cat((out,sp),1)
|
||||
|
||||
|
||||
out = self.conv3(out)
|
||||
out = self.bn3(out)
|
||||
|
||||
residual = self.shortcut(x)
|
||||
out += residual
|
||||
out = self.relu(out)
|
||||
|
||||
return out
|
||||
|
||||
class ERes2Net(nn.Module):
|
||||
def __init__(self,
|
||||
block=BasicBlockERes2Net,
|
||||
block_fuse=BasicBlockERes2Net_diff_AFF,
|
||||
num_blocks=[3, 4, 6, 3],
|
||||
m_channels=64,
|
||||
feat_dim=80,
|
||||
embedding_size=192,
|
||||
pooling_func='TSTP',
|
||||
two_emb_layer=False):
|
||||
super(ERes2Net, self).__init__()
|
||||
self.in_planes = m_channels
|
||||
self.feat_dim = feat_dim
|
||||
self.embedding_size = embedding_size
|
||||
self.stats_dim = int(feat_dim / 8) * m_channels * 8
|
||||
self.two_emb_layer = two_emb_layer
|
||||
|
||||
self.conv1 = nn.Conv2d(1, m_channels, kernel_size=3, stride=1, padding=1, bias=False)
|
||||
self.bn1 = nn.BatchNorm2d(m_channels)
|
||||
|
||||
self.layer1 = self._make_layer(block, m_channels, num_blocks[0], stride=1)
|
||||
self.layer2 = self._make_layer(block, m_channels * 2, num_blocks[1], stride=2)
|
||||
self.layer3 = self._make_layer(block_fuse, m_channels * 4, num_blocks[2], stride=2)
|
||||
self.layer4 = self._make_layer(block_fuse, m_channels * 8, num_blocks[3], stride=2)
|
||||
|
||||
self.layer1_downsample = nn.Conv2d(m_channels * 4, m_channels * 8, kernel_size=3, padding=1, stride=2, bias=False)
|
||||
self.layer2_downsample = nn.Conv2d(m_channels * 8, m_channels * 16, kernel_size=3, padding=1, stride=2, bias=False)
|
||||
self.layer3_downsample = nn.Conv2d(m_channels * 16, m_channels * 32, kernel_size=3, padding=1, stride=2, bias=False)
|
||||
|
||||
self.fuse_mode12 = AFF(channels=m_channels * 8)
|
||||
self.fuse_mode123 = AFF(channels=m_channels * 16)
|
||||
self.fuse_mode1234 = AFF(channels=m_channels * 32)
|
||||
|
||||
self.n_stats = 1 if pooling_func == 'TAP' or pooling_func == "TSDP" else 2
|
||||
self.pool = getattr(pooling_layers, pooling_func)(
|
||||
in_dim=self.stats_dim * block.expansion)
|
||||
self.seg_1 = nn.Linear(self.stats_dim * block.expansion * self.n_stats, embedding_size)
|
||||
if self.two_emb_layer:
|
||||
self.seg_bn_1 = nn.BatchNorm1d(embedding_size, affine=False)
|
||||
self.seg_2 = nn.Linear(embedding_size, embedding_size)
|
||||
else:
|
||||
self.seg_bn_1 = nn.Identity()
|
||||
self.seg_2 = nn.Identity()
|
||||
|
||||
def _make_layer(self, block, planes, num_blocks, stride):
|
||||
strides = [stride] + [1] * (num_blocks - 1)
|
||||
layers = []
|
||||
for stride in strides:
|
||||
layers.append(block(self.in_planes, planes, stride))
|
||||
self.in_planes = planes * block.expansion
|
||||
return nn.Sequential(*layers)
|
||||
|
||||
def forward(self, x):
|
||||
x = x.permute(0, 2, 1) # (B,T,F) => (B,F,T)
|
||||
|
||||
x = x.unsqueeze_(1)
|
||||
out = F.relu(self.bn1(self.conv1(x)))
|
||||
out1 = self.layer1(out)
|
||||
out2 = self.layer2(out1)
|
||||
out1_downsample = self.layer1_downsample(out1)
|
||||
fuse_out12 = self.fuse_mode12(out2, out1_downsample)
|
||||
out3 = self.layer3(out2)
|
||||
fuse_out12_downsample = self.layer2_downsample(fuse_out12)
|
||||
fuse_out123 = self.fuse_mode123(out3, fuse_out12_downsample)
|
||||
out4 = self.layer4(out3)
|
||||
fuse_out123_downsample = self.layer3_downsample(fuse_out123)
|
||||
fuse_out1234 = self.fuse_mode1234(out4, fuse_out123_downsample)
|
||||
stats = self.pool(fuse_out1234)
|
||||
|
||||
embed_a = self.seg_1(stats)
|
||||
if self.two_emb_layer:
|
||||
out = F.relu(embed_a)
|
||||
out = self.seg_bn_1(out)
|
||||
embed_b = self.seg_2(out)
|
||||
return embed_b
|
||||
else:
|
||||
return embed_a
|
||||
|
||||
def forward2(self, x,if_mean):
|
||||
x = x.permute(0, 2, 1) # (B,T,F) => (B,F,T)
|
||||
|
||||
x = x.unsqueeze_(1)
|
||||
out = F.relu(self.bn1(self.conv1(x)))
|
||||
out1 = self.layer1(out)
|
||||
out2 = self.layer2(out1)
|
||||
out1_downsample = self.layer1_downsample(out1)
|
||||
fuse_out12 = self.fuse_mode12(out2, out1_downsample)
|
||||
out3 = self.layer3(out2)
|
||||
fuse_out12_downsample = self.layer2_downsample(fuse_out12)
|
||||
fuse_out123 = self.fuse_mode123(out3, fuse_out12_downsample)
|
||||
out4 = self.layer4(out3)
|
||||
fuse_out123_downsample = self.layer3_downsample(fuse_out123)
|
||||
fuse_out1234 = self.fuse_mode1234(out4, fuse_out123_downsample).flatten(start_dim=1,end_dim=2)#bs,20480,T
|
||||
if(if_mean==False):
|
||||
mean=fuse_out1234[0].transpose(1,0)#(T,20480),bs=T
|
||||
else:
|
||||
mean = fuse_out1234.mean(2)#bs,20480
|
||||
mean_std=torch.cat([mean,torch.zeros_like(mean)],1)
|
||||
return self.seg_1(mean_std)#(T,192)
|
||||
|
||||
|
||||
# stats = self.pool(fuse_out1234)
|
||||
# if self.two_emb_layer:
|
||||
# out = F.relu(embed_a)
|
||||
# out = self.seg_bn_1(out)
|
||||
# embed_b = self.seg_2(out)
|
||||
# return embed_b
|
||||
# else:
|
||||
# return embed_a
|
||||
|
||||
def forward3(self, x):
|
||||
x = x.permute(0, 2, 1) # (B,T,F) => (B,F,T)
|
||||
|
||||
x = x.unsqueeze_(1)
|
||||
out = F.relu(self.bn1(self.conv1(x)))
|
||||
out1 = self.layer1(out)
|
||||
out2 = self.layer2(out1)
|
||||
out1_downsample = self.layer1_downsample(out1)
|
||||
fuse_out12 = self.fuse_mode12(out2, out1_downsample)
|
||||
out3 = self.layer3(out2)
|
||||
fuse_out12_downsample = self.layer2_downsample(fuse_out12)
|
||||
fuse_out123 = self.fuse_mode123(out3, fuse_out12_downsample)
|
||||
out4 = self.layer4(out3)
|
||||
fuse_out123_downsample = self.layer3_downsample(fuse_out123)
|
||||
fuse_out1234 = self.fuse_mode1234(out4, fuse_out123_downsample).flatten(start_dim=1,end_dim=2).mean(-1)
|
||||
return fuse_out1234
|
||||
# print(fuse_out1234.shape)
|
||||
# print(fuse_out1234.flatten(start_dim=1,end_dim=2).shape)
|
||||
# pdb.set_trace()
|
||||
|
||||
|
||||
|
||||
|
||||
@@ -6,7 +6,6 @@ import torch.nn as nn
|
||||
|
||||
|
||||
class AFF(nn.Module):
|
||||
|
||||
def __init__(self, channels=64, r=4):
|
||||
super(AFF, self).__init__()
|
||||
inter_channels = int(channels // r)
|
||||
@@ -23,7 +22,6 @@ class AFF(nn.Module):
|
||||
xa = torch.cat((x, ds_y), dim=1)
|
||||
x_att = self.local_att(xa)
|
||||
x_att = 1.0 + torch.tanh(x_att)
|
||||
xo = torch.mul(x, x_att) + torch.mul(ds_y, 2.0-x_att)
|
||||
xo = torch.mul(x, x_att) + torch.mul(ds_y, 2.0 - x_att)
|
||||
|
||||
return xo
|
||||
|
||||
|
||||
+37
-277
@@ -2,18 +2,14 @@ import math
|
||||
from typing import Tuple
|
||||
|
||||
import torch
|
||||
import torchaudio
|
||||
from torch import Tensor
|
||||
|
||||
__all__ = [
|
||||
"get_mel_banks",
|
||||
"inverse_mel_scale",
|
||||
"inverse_mel_scale_scalar",
|
||||
"mel_scale",
|
||||
"mel_scale_scalar",
|
||||
"spectrogram",
|
||||
"fbank",
|
||||
"mfcc",
|
||||
"vtln_warp_freq",
|
||||
"vtln_warp_mel_freq",
|
||||
]
|
||||
@@ -144,7 +140,7 @@ def _get_waveform_and_window_properties(
|
||||
)
|
||||
assert 0 < window_shift, "`window_shift` must be greater than 0"
|
||||
assert padded_window_size % 2 == 0, (
|
||||
"the padded `window_size` must be divisible by two." " use `round_to_power_of_two` or change `frame_length`"
|
||||
"the padded `window_size` must be divisible by two. use `round_to_power_of_two` or change `frame_length`"
|
||||
)
|
||||
assert 0.0 <= preemphasis_coefficient <= 1.0, "`preemphasis_coefficient` must be between [0,1]"
|
||||
assert sample_frequency > 0, "`sample_frequency` must be greater than zero"
|
||||
@@ -226,99 +222,6 @@ def _subtract_column_mean(tensor: Tensor, subtract_mean: bool) -> Tensor:
|
||||
return tensor
|
||||
|
||||
|
||||
def spectrogram(
|
||||
waveform: Tensor,
|
||||
blackman_coeff: float = 0.42,
|
||||
channel: int = -1,
|
||||
dither: float = 0.0,
|
||||
energy_floor: float = 1.0,
|
||||
frame_length: float = 25.0,
|
||||
frame_shift: float = 10.0,
|
||||
min_duration: float = 0.0,
|
||||
preemphasis_coefficient: float = 0.97,
|
||||
raw_energy: bool = True,
|
||||
remove_dc_offset: bool = True,
|
||||
round_to_power_of_two: bool = True,
|
||||
sample_frequency: float = 16000.0,
|
||||
snip_edges: bool = True,
|
||||
subtract_mean: bool = False,
|
||||
window_type: str = POVEY,
|
||||
) -> Tensor:
|
||||
r"""Create a spectrogram from a raw audio signal. This matches the input/output of Kaldi's
|
||||
compute-spectrogram-feats.
|
||||
|
||||
Args:
|
||||
waveform (Tensor): Tensor of audio of size (c, n) where c is in the range [0,2)
|
||||
blackman_coeff (float, optional): Constant coefficient for generalized Blackman window. (Default: ``0.42``)
|
||||
channel (int, optional): Channel to extract (-1 -> expect mono, 0 -> left, 1 -> right) (Default: ``-1``)
|
||||
dither (float, optional): Dithering constant (0.0 means no dither). If you turn this off, you should set
|
||||
the energy_floor option, e.g. to 1.0 or 0.1 (Default: ``0.0``)
|
||||
energy_floor (float, optional): Floor on energy (absolute, not relative) in Spectrogram computation. Caution:
|
||||
this floor is applied to the zeroth component, representing the total signal energy. The floor on the
|
||||
individual spectrogram elements is fixed at std::numeric_limits<float>::epsilon(). (Default: ``1.0``)
|
||||
frame_length (float, optional): Frame length in milliseconds (Default: ``25.0``)
|
||||
frame_shift (float, optional): Frame shift in milliseconds (Default: ``10.0``)
|
||||
min_duration (float, optional): Minimum duration of segments to process (in seconds). (Default: ``0.0``)
|
||||
preemphasis_coefficient (float, optional): Coefficient for use in signal preemphasis (Default: ``0.97``)
|
||||
raw_energy (bool, optional): If True, compute energy before preemphasis and windowing (Default: ``True``)
|
||||
remove_dc_offset (bool, optional): Subtract mean from waveform on each frame (Default: ``True``)
|
||||
round_to_power_of_two (bool, optional): If True, round window size to power of two by zero-padding input
|
||||
to FFT. (Default: ``True``)
|
||||
sample_frequency (float, optional): Waveform data sample frequency (must match the waveform file, if
|
||||
specified there) (Default: ``16000.0``)
|
||||
snip_edges (bool, optional): If True, end effects will be handled by outputting only frames that completely fit
|
||||
in the file, and the number of frames depends on the frame_length. If False, the number of frames
|
||||
depends only on the frame_shift, and we reflect the data at the ends. (Default: ``True``)
|
||||
subtract_mean (bool, optional): Subtract mean of each feature file [CMS]; not recommended to do
|
||||
it this way. (Default: ``False``)
|
||||
window_type (str, optional): Type of window ('hamming'|'hanning'|'povey'|'rectangular'|'blackman')
|
||||
(Default: ``'povey'``)
|
||||
|
||||
Returns:
|
||||
Tensor: A spectrogram identical to what Kaldi would output. The shape is
|
||||
(m, ``padded_window_size // 2 + 1``) where m is calculated in _get_strided
|
||||
"""
|
||||
device, dtype = waveform.device, waveform.dtype
|
||||
epsilon = _get_epsilon(device, dtype)
|
||||
|
||||
waveform, window_shift, window_size, padded_window_size = _get_waveform_and_window_properties(
|
||||
waveform, channel, sample_frequency, frame_shift, frame_length, round_to_power_of_two, preemphasis_coefficient
|
||||
)
|
||||
|
||||
if len(waveform) < min_duration * sample_frequency:
|
||||
# signal is too short
|
||||
return torch.empty(0)
|
||||
|
||||
strided_input, signal_log_energy = _get_window(
|
||||
waveform,
|
||||
padded_window_size,
|
||||
window_size,
|
||||
window_shift,
|
||||
window_type,
|
||||
blackman_coeff,
|
||||
snip_edges,
|
||||
raw_energy,
|
||||
energy_floor,
|
||||
dither,
|
||||
remove_dc_offset,
|
||||
preemphasis_coefficient,
|
||||
)
|
||||
|
||||
# size (m, padded_window_size // 2 + 1, 2)
|
||||
fft = torch.fft.rfft(strided_input)
|
||||
|
||||
# Convert the FFT into a power spectrum
|
||||
power_spectrum = torch.max(fft.abs().pow(2.0), epsilon).log() # size (m, padded_window_size // 2 + 1)
|
||||
power_spectrum[:, 0] = signal_log_energy
|
||||
|
||||
power_spectrum = _subtract_column_mean(power_spectrum, subtract_mean)
|
||||
return power_spectrum
|
||||
|
||||
|
||||
def inverse_mel_scale_scalar(mel_freq: float) -> float:
|
||||
return 700.0 * (math.exp(mel_freq / 1127.0) - 1.0)
|
||||
|
||||
|
||||
def inverse_mel_scale(mel_freq: Tensor) -> Tensor:
|
||||
return 700.0 * ((mel_freq / 1127.0).exp() - 1.0)
|
||||
|
||||
@@ -441,7 +344,9 @@ def get_mel_banks(
|
||||
high_freq: float,
|
||||
vtln_low: float,
|
||||
vtln_high: float,
|
||||
vtln_warp_factor: float,device=None,dtype=None
|
||||
vtln_warp_factor: float,
|
||||
device=None,
|
||||
dtype=None,
|
||||
) -> Tuple[Tensor, Tensor]:
|
||||
"""
|
||||
Returns:
|
||||
@@ -457,9 +362,9 @@ def get_mel_banks(
|
||||
if high_freq <= 0.0:
|
||||
high_freq += nyquist
|
||||
|
||||
assert (
|
||||
(0.0 <= low_freq < nyquist) and (0.0 < high_freq <= nyquist) and (low_freq < high_freq)
|
||||
), "Bad values in options: low-freq {} and high-freq {} vs. nyquist {}".format(low_freq, high_freq, nyquist)
|
||||
assert (0.0 <= low_freq < nyquist) and (0.0 < high_freq <= nyquist) and (low_freq < high_freq), (
|
||||
"Bad values in options: low-freq {} and high-freq {} vs. nyquist {}".format(low_freq, high_freq, nyquist)
|
||||
)
|
||||
|
||||
# fft-bin width [think of it as Nyquist-freq / half-window-length]
|
||||
fft_bin_width = sample_freq / window_length_padded
|
||||
@@ -475,7 +380,7 @@ def get_mel_banks(
|
||||
|
||||
assert vtln_warp_factor == 1.0 or (
|
||||
(low_freq < vtln_low < high_freq) and (0.0 < vtln_high < high_freq) and (vtln_low < vtln_high)
|
||||
), "Bad values in options: vtln-low {} and vtln-high {}, versus " "low-freq {} and high-freq {}".format(
|
||||
), "Bad values in options: vtln-low {} and vtln-high {}, versus low-freq {} and high-freq {}".format(
|
||||
vtln_low, vtln_high, low_freq, high_freq
|
||||
)
|
||||
|
||||
@@ -508,9 +413,12 @@ def get_mel_banks(
|
||||
bins[up_idx] = up_slope[up_idx]
|
||||
bins[down_idx] = down_slope[down_idx]
|
||||
|
||||
return bins.to(device=device,dtype=dtype)#, center_freqs
|
||||
return bins.to(device=device, dtype=dtype) # , center_freqs
|
||||
|
||||
|
||||
cache = {}
|
||||
|
||||
|
||||
cache={}
|
||||
def fbank(
|
||||
waveform: Tensor,
|
||||
blackman_coeff: float = 0.42,
|
||||
@@ -620,14 +528,34 @@ def fbank(
|
||||
# size (num_mel_bins, padded_window_size // 2)
|
||||
# print(num_mel_bins, padded_window_size, sample_frequency, low_freq, high_freq, vtln_low, vtln_high, vtln_warp)
|
||||
|
||||
cache_key="%s-%s-%s-%s-%s-%s-%s-%s-%s-%s"%(num_mel_bins, padded_window_size, sample_frequency, low_freq, high_freq, vtln_low, vtln_high, vtln_warp,device,dtype)
|
||||
cache_key = "%s-%s-%s-%s-%s-%s-%s-%s-%s-%s" % (
|
||||
num_mel_bins,
|
||||
padded_window_size,
|
||||
sample_frequency,
|
||||
low_freq,
|
||||
high_freq,
|
||||
vtln_low,
|
||||
vtln_high,
|
||||
vtln_warp,
|
||||
device,
|
||||
dtype,
|
||||
)
|
||||
if cache_key not in cache:
|
||||
mel_energies = get_mel_banks(
|
||||
num_mel_bins, padded_window_size, sample_frequency, low_freq, high_freq, vtln_low, vtln_high, vtln_warp,device,dtype
|
||||
num_mel_bins,
|
||||
padded_window_size,
|
||||
sample_frequency,
|
||||
low_freq,
|
||||
high_freq,
|
||||
vtln_low,
|
||||
vtln_high,
|
||||
vtln_warp,
|
||||
device,
|
||||
dtype,
|
||||
)
|
||||
cache[cache_key]=mel_energies
|
||||
cache[cache_key] = mel_energies
|
||||
else:
|
||||
mel_energies=cache[cache_key]
|
||||
mel_energies = cache[cache_key]
|
||||
|
||||
# pad right column with zeros and add dimension, size (num_mel_bins, padded_window_size // 2 + 1)
|
||||
mel_energies = torch.nn.functional.pad(mel_energies, (0, 1), mode="constant", value=0)
|
||||
@@ -649,171 +577,3 @@ def fbank(
|
||||
|
||||
mel_energies = _subtract_column_mean(mel_energies, subtract_mean)
|
||||
return mel_energies
|
||||
|
||||
|
||||
def _get_dct_matrix(num_ceps: int, num_mel_bins: int) -> Tensor:
|
||||
# returns a dct matrix of size (num_mel_bins, num_ceps)
|
||||
# size (num_mel_bins, num_mel_bins)
|
||||
dct_matrix = torchaudio.functional.create_dct(num_mel_bins, num_mel_bins, "ortho")
|
||||
# kaldi expects the first cepstral to be weighted sum of factor sqrt(1/num_mel_bins)
|
||||
# this would be the first column in the dct_matrix for torchaudio as it expects a
|
||||
# right multiply (which would be the first column of the kaldi's dct_matrix as kaldi
|
||||
# expects a left multiply e.g. dct_matrix * vector).
|
||||
dct_matrix[:, 0] = math.sqrt(1 / float(num_mel_bins))
|
||||
dct_matrix = dct_matrix[:, :num_ceps]
|
||||
return dct_matrix
|
||||
|
||||
|
||||
def _get_lifter_coeffs(num_ceps: int, cepstral_lifter: float) -> Tensor:
|
||||
# returns size (num_ceps)
|
||||
# Compute liftering coefficients (scaling on cepstral coeffs)
|
||||
# coeffs are numbered slightly differently from HTK: the zeroth index is C0, which is not affected.
|
||||
i = torch.arange(num_ceps)
|
||||
return 1.0 + 0.5 * cepstral_lifter * torch.sin(math.pi * i / cepstral_lifter)
|
||||
|
||||
|
||||
def mfcc(
|
||||
waveform: Tensor,
|
||||
blackman_coeff: float = 0.42,
|
||||
cepstral_lifter: float = 22.0,
|
||||
channel: int = -1,
|
||||
dither: float = 0.0,
|
||||
energy_floor: float = 1.0,
|
||||
frame_length: float = 25.0,
|
||||
frame_shift: float = 10.0,
|
||||
high_freq: float = 0.0,
|
||||
htk_compat: bool = False,
|
||||
low_freq: float = 20.0,
|
||||
num_ceps: int = 13,
|
||||
min_duration: float = 0.0,
|
||||
num_mel_bins: int = 23,
|
||||
preemphasis_coefficient: float = 0.97,
|
||||
raw_energy: bool = True,
|
||||
remove_dc_offset: bool = True,
|
||||
round_to_power_of_two: bool = True,
|
||||
sample_frequency: float = 16000.0,
|
||||
snip_edges: bool = True,
|
||||
subtract_mean: bool = False,
|
||||
use_energy: bool = False,
|
||||
vtln_high: float = -500.0,
|
||||
vtln_low: float = 100.0,
|
||||
vtln_warp: float = 1.0,
|
||||
window_type: str = POVEY,
|
||||
) -> Tensor:
|
||||
r"""Create a mfcc from a raw audio signal. This matches the input/output of Kaldi's
|
||||
compute-mfcc-feats.
|
||||
|
||||
Args:
|
||||
waveform (Tensor): Tensor of audio of size (c, n) where c is in the range [0,2)
|
||||
blackman_coeff (float, optional): Constant coefficient for generalized Blackman window. (Default: ``0.42``)
|
||||
cepstral_lifter (float, optional): Constant that controls scaling of MFCCs (Default: ``22.0``)
|
||||
channel (int, optional): Channel to extract (-1 -> expect mono, 0 -> left, 1 -> right) (Default: ``-1``)
|
||||
dither (float, optional): Dithering constant (0.0 means no dither). If you turn this off, you should set
|
||||
the energy_floor option, e.g. to 1.0 or 0.1 (Default: ``0.0``)
|
||||
energy_floor (float, optional): Floor on energy (absolute, not relative) in Spectrogram computation. Caution:
|
||||
this floor is applied to the zeroth component, representing the total signal energy. The floor on the
|
||||
individual spectrogram elements is fixed at std::numeric_limits<float>::epsilon(). (Default: ``1.0``)
|
||||
frame_length (float, optional): Frame length in milliseconds (Default: ``25.0``)
|
||||
frame_shift (float, optional): Frame shift in milliseconds (Default: ``10.0``)
|
||||
high_freq (float, optional): High cutoff frequency for mel bins (if <= 0, offset from Nyquist)
|
||||
(Default: ``0.0``)
|
||||
htk_compat (bool, optional): If true, put energy last. Warning: not sufficient to get HTK compatible
|
||||
features (need to change other parameters). (Default: ``False``)
|
||||
low_freq (float, optional): Low cutoff frequency for mel bins (Default: ``20.0``)
|
||||
num_ceps (int, optional): Number of cepstra in MFCC computation (including C0) (Default: ``13``)
|
||||
min_duration (float, optional): Minimum duration of segments to process (in seconds). (Default: ``0.0``)
|
||||
num_mel_bins (int, optional): Number of triangular mel-frequency bins (Default: ``23``)
|
||||
preemphasis_coefficient (float, optional): Coefficient for use in signal preemphasis (Default: ``0.97``)
|
||||
raw_energy (bool, optional): If True, compute energy before preemphasis and windowing (Default: ``True``)
|
||||
remove_dc_offset (bool, optional): Subtract mean from waveform on each frame (Default: ``True``)
|
||||
round_to_power_of_two (bool, optional): If True, round window size to power of two by zero-padding input
|
||||
to FFT. (Default: ``True``)
|
||||
sample_frequency (float, optional): Waveform data sample frequency (must match the waveform file, if
|
||||
specified there) (Default: ``16000.0``)
|
||||
snip_edges (bool, optional): If True, end effects will be handled by outputting only frames that completely fit
|
||||
in the file, and the number of frames depends on the frame_length. If False, the number of frames
|
||||
depends only on the frame_shift, and we reflect the data at the ends. (Default: ``True``)
|
||||
subtract_mean (bool, optional): Subtract mean of each feature file [CMS]; not recommended to do
|
||||
it this way. (Default: ``False``)
|
||||
use_energy (bool, optional): Add an extra dimension with energy to the FBANK output. (Default: ``False``)
|
||||
vtln_high (float, optional): High inflection point in piecewise linear VTLN warping function (if
|
||||
negative, offset from high-mel-freq (Default: ``-500.0``)
|
||||
vtln_low (float, optional): Low inflection point in piecewise linear VTLN warping function (Default: ``100.0``)
|
||||
vtln_warp (float, optional): Vtln warp factor (only applicable if vtln_map not specified) (Default: ``1.0``)
|
||||
window_type (str, optional): Type of window ('hamming'|'hanning'|'povey'|'rectangular'|'blackman')
|
||||
(Default: ``"povey"``)
|
||||
|
||||
Returns:
|
||||
Tensor: A mfcc identical to what Kaldi would output. The shape is (m, ``num_ceps``)
|
||||
where m is calculated in _get_strided
|
||||
"""
|
||||
assert num_ceps <= num_mel_bins, "num_ceps cannot be larger than num_mel_bins: %d vs %d" % (num_ceps, num_mel_bins)
|
||||
|
||||
device, dtype = waveform.device, waveform.dtype
|
||||
|
||||
# The mel_energies should not be squared (use_power=True), not have mean subtracted
|
||||
# (subtract_mean=False), and use log (use_log_fbank=True).
|
||||
# size (m, num_mel_bins + use_energy)
|
||||
feature = fbank(
|
||||
waveform=waveform,
|
||||
blackman_coeff=blackman_coeff,
|
||||
channel=channel,
|
||||
dither=dither,
|
||||
energy_floor=energy_floor,
|
||||
frame_length=frame_length,
|
||||
frame_shift=frame_shift,
|
||||
high_freq=high_freq,
|
||||
htk_compat=htk_compat,
|
||||
low_freq=low_freq,
|
||||
min_duration=min_duration,
|
||||
num_mel_bins=num_mel_bins,
|
||||
preemphasis_coefficient=preemphasis_coefficient,
|
||||
raw_energy=raw_energy,
|
||||
remove_dc_offset=remove_dc_offset,
|
||||
round_to_power_of_two=round_to_power_of_two,
|
||||
sample_frequency=sample_frequency,
|
||||
snip_edges=snip_edges,
|
||||
subtract_mean=False,
|
||||
use_energy=use_energy,
|
||||
use_log_fbank=True,
|
||||
use_power=True,
|
||||
vtln_high=vtln_high,
|
||||
vtln_low=vtln_low,
|
||||
vtln_warp=vtln_warp,
|
||||
window_type=window_type,
|
||||
)
|
||||
|
||||
if use_energy:
|
||||
# size (m)
|
||||
signal_log_energy = feature[:, num_mel_bins if htk_compat else 0]
|
||||
# offset is 0 if htk_compat==True else 1
|
||||
mel_offset = int(not htk_compat)
|
||||
feature = feature[:, mel_offset : (num_mel_bins + mel_offset)]
|
||||
|
||||
# size (num_mel_bins, num_ceps)
|
||||
dct_matrix = _get_dct_matrix(num_ceps, num_mel_bins).to(dtype=dtype, device=device)
|
||||
|
||||
# size (m, num_ceps)
|
||||
feature = feature.matmul(dct_matrix)
|
||||
|
||||
if cepstral_lifter != 0.0:
|
||||
# size (1, num_ceps)
|
||||
lifter_coeffs = _get_lifter_coeffs(num_ceps, cepstral_lifter).unsqueeze(0)
|
||||
feature *= lifter_coeffs.to(device=device, dtype=dtype)
|
||||
|
||||
# if use_energy then replace the last column for htk_compat == true else first column
|
||||
if use_energy:
|
||||
feature[:, 0] = signal_log_energy
|
||||
|
||||
if htk_compat:
|
||||
energy = feature[:, 0].unsqueeze(1) # size (m, 1)
|
||||
feature = feature[:, 1:] # size (m, num_ceps - 1)
|
||||
if not use_energy:
|
||||
# scale on C0 (actually removing a scale we previously added that's
|
||||
# part of one common definition of the cosine transform.)
|
||||
energy *= math.sqrt(2)
|
||||
|
||||
feature = torch.cat((feature, energy), dim=1)
|
||||
|
||||
feature = _subtract_column_mean(feature, subtract_mean)
|
||||
return feature
|
||||
|
||||
@@ -1,7 +1,7 @@
|
||||
# Copyright 3D-Speaker (https://github.com/alibaba-damo-academy/3D-Speaker). All Rights Reserved.
|
||||
# Licensed under the Apache License, Version 2.0 (http://www.apache.org/licenses/LICENSE-2.0)
|
||||
|
||||
""" This implementation is adapted from https://github.com/wenet-e2e/wespeaker."""
|
||||
"""This implementation is adapted from https://github.com/wenet-e2e/wespeaker."""
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
@@ -11,6 +11,7 @@ class TAP(nn.Module):
|
||||
"""
|
||||
Temporal average pooling, only first-order mean is considered
|
||||
"""
|
||||
|
||||
def __init__(self, **kwargs):
|
||||
super(TAP, self).__init__()
|
||||
|
||||
@@ -25,6 +26,7 @@ class TSDP(nn.Module):
|
||||
"""
|
||||
Temporal standard deviation pooling, only second-order std is considered
|
||||
"""
|
||||
|
||||
def __init__(self, **kwargs):
|
||||
super(TSDP, self).__init__()
|
||||
|
||||
@@ -41,6 +43,7 @@ class TSTP(nn.Module):
|
||||
x-vector
|
||||
Comment: simple concatenation can not make full use of both statistics
|
||||
"""
|
||||
|
||||
def __init__(self, **kwargs):
|
||||
super(TSTP, self).__init__()
|
||||
|
||||
@@ -56,9 +59,10 @@ class TSTP(nn.Module):
|
||||
|
||||
|
||||
class ASTP(nn.Module):
|
||||
""" Attentive statistics pooling: Channel- and context-dependent
|
||||
statistics pooling, first used in ECAPA_TDNN.
|
||||
"""Attentive statistics pooling: Channel- and context-dependent
|
||||
statistics pooling, first used in ECAPA_TDNN.
|
||||
"""
|
||||
|
||||
def __init__(self, in_dim, bottleneck_dim=128, global_context_att=False):
|
||||
super(ASTP, self).__init__()
|
||||
self.global_context_att = global_context_att
|
||||
@@ -66,15 +70,10 @@ class ASTP(nn.Module):
|
||||
# Use Conv1d with stride == 1 rather than Linear, then we don't
|
||||
# need to transpose inputs.
|
||||
if global_context_att:
|
||||
self.linear1 = nn.Conv1d(
|
||||
in_dim * 3, bottleneck_dim,
|
||||
kernel_size=1) # equals W and b in the paper
|
||||
self.linear1 = nn.Conv1d(in_dim * 3, bottleneck_dim, kernel_size=1) # equals W and b in the paper
|
||||
else:
|
||||
self.linear1 = nn.Conv1d(
|
||||
in_dim, bottleneck_dim,
|
||||
kernel_size=1) # equals W and b in the paper
|
||||
self.linear2 = nn.Conv1d(bottleneck_dim, in_dim,
|
||||
kernel_size=1) # equals V and k in the paper
|
||||
self.linear1 = nn.Conv1d(in_dim, bottleneck_dim, kernel_size=1) # equals W and b in the paper
|
||||
self.linear2 = nn.Conv1d(bottleneck_dim, in_dim, kernel_size=1) # equals V and k in the paper
|
||||
|
||||
def forward(self, x):
|
||||
"""
|
||||
@@ -88,15 +87,13 @@ class ASTP(nn.Module):
|
||||
|
||||
if self.global_context_att:
|
||||
context_mean = torch.mean(x, dim=-1, keepdim=True).expand_as(x)
|
||||
context_std = torch.sqrt(
|
||||
torch.var(x, dim=-1, keepdim=True) + 1e-10).expand_as(x)
|
||||
context_std = torch.sqrt(torch.var(x, dim=-1, keepdim=True) + 1e-10).expand_as(x)
|
||||
x_in = torch.cat((x, context_mean, context_std), dim=1)
|
||||
else:
|
||||
x_in = x
|
||||
|
||||
# DON'T use ReLU here! ReLU may be hard to converge.
|
||||
alpha = torch.tanh(
|
||||
self.linear1(x_in)) # alpha = F.relu(self.linear1(x_in))
|
||||
alpha = torch.tanh(self.linear1(x_in)) # alpha = F.relu(self.linear1(x_in))
|
||||
alpha = torch.softmax(self.linear2(alpha), dim=2)
|
||||
mean = torch.sum(alpha * x, dim=2)
|
||||
var = torch.sum(alpha * (x**2), dim=2) - mean**2
|
||||
|
||||
@@ -1,6 +1,7 @@
|
||||
# modified from https://github.com/yangdongchao/SoundStorm/blob/master/soundstorm/s1/AR/models/t2s_model.py
|
||||
# reference: https://github.com/lifeiteng/vall-e
|
||||
import argparse
|
||||
from io import BytesIO
|
||||
from typing import Optional
|
||||
from my_utils import load_audio
|
||||
import torch
|
||||
@@ -17,6 +18,9 @@ from module.models_onnx import SynthesizerTrn
|
||||
|
||||
from inference_webui import get_phones_and_bert
|
||||
|
||||
from sv import SV
|
||||
import kaldi as Kaldi
|
||||
|
||||
import os
|
||||
import soundfile
|
||||
|
||||
@@ -32,6 +36,25 @@ default_config = {
|
||||
"EOS": 1024,
|
||||
}
|
||||
|
||||
sv_cn_model = None
|
||||
|
||||
|
||||
def init_sv_cn(device, is_half):
|
||||
global sv_cn_model
|
||||
sv_cn_model = SV(device, is_half)
|
||||
|
||||
|
||||
def load_sovits_new(sovits_path):
|
||||
f = open(sovits_path, "rb")
|
||||
meta = f.read(2)
|
||||
if meta != b"PK":
|
||||
data = b"PK" + f.read()
|
||||
bio = BytesIO()
|
||||
bio.write(data)
|
||||
bio.seek(0)
|
||||
return torch.load(bio, map_location="cpu", weights_only=False)
|
||||
return torch.load(sovits_path, map_location="cpu", weights_only=False)
|
||||
|
||||
|
||||
def get_raw_t2s_model(dict_s1) -> Text2SemanticLightningModule:
|
||||
config = dict_s1["config"]
|
||||
@@ -83,7 +106,7 @@ def logits_to_probs(
|
||||
@torch.jit.script
|
||||
def multinomial_sample_one_no_sync(probs_sort):
|
||||
# Does multinomial sampling without a cuda synchronization
|
||||
q = torch.randn_like(probs_sort)
|
||||
q = torch.empty_like(probs_sort).exponential_(1.0)
|
||||
return torch.argmax(probs_sort / q, dim=-1, keepdim=True).to(dtype=torch.int)
|
||||
|
||||
|
||||
@@ -94,7 +117,7 @@ def sample(
|
||||
temperature: float = 1.0,
|
||||
top_k: Optional[int] = None,
|
||||
top_p: Optional[int] = None,
|
||||
repetition_penalty: float = 1.0,
|
||||
repetition_penalty: float = 1.35,
|
||||
):
|
||||
probs = logits_to_probs(
|
||||
logits=logits,
|
||||
@@ -109,8 +132,10 @@ def sample(
|
||||
|
||||
|
||||
@torch.jit.script
|
||||
def spectrogram_torch(y: Tensor, n_fft: int, sampling_rate: int, hop_size: int, win_size: int, center: bool = False):
|
||||
hann_window = torch.hann_window(win_size, device=y.device, dtype=y.dtype)
|
||||
def spectrogram_torch(
|
||||
hann_window: Tensor, y: Tensor, n_fft: int, sampling_rate: int, hop_size: int, win_size: int, center: bool = False
|
||||
):
|
||||
# hann_window = torch.hann_window(win_size, device=y.device, dtype=y.dtype)
|
||||
y = torch.nn.functional.pad(
|
||||
y.unsqueeze(1),
|
||||
(int((n_fft - hop_size) / 2), int((n_fft - hop_size) / 2)),
|
||||
@@ -236,41 +261,21 @@ class T2SBlock:
|
||||
|
||||
attn = F.scaled_dot_product_attention(q, k, v, ~attn_mask)
|
||||
|
||||
attn = attn.permute(2, 0, 1, 3).reshape(batch_size * q_len, self.hidden_dim)
|
||||
attn = attn.view(q_len, batch_size, self.hidden_dim).transpose(1, 0)
|
||||
# attn = attn.permute(2, 0, 1, 3).reshape(batch_size * q_len, self.hidden_dim)
|
||||
# attn = attn.view(q_len, batch_size, self.hidden_dim).transpose(1, 0)
|
||||
attn = attn.transpose(1, 2).reshape(batch_size, q_len, -1)
|
||||
attn = F.linear(self.to_mask(attn, padding_mask), self.out_w, self.out_b)
|
||||
|
||||
if padding_mask is not None:
|
||||
for i in range(batch_size):
|
||||
# mask = padding_mask[i,:,0]
|
||||
if self.false.device != padding_mask.device:
|
||||
self.false = self.false.to(padding_mask.device)
|
||||
idx = torch.where(padding_mask[i, :, 0] == self.false)[0]
|
||||
x_item = x[i, idx, :].unsqueeze(0)
|
||||
attn_item = attn[i, idx, :].unsqueeze(0)
|
||||
x_item = x_item + attn_item
|
||||
x_item = F.layer_norm(x_item, [self.hidden_dim], self.norm_w1, self.norm_b1, self.norm_eps1)
|
||||
x_item = x_item + self.mlp.forward(x_item)
|
||||
x_item = F.layer_norm(
|
||||
x_item,
|
||||
[self.hidden_dim],
|
||||
self.norm_w2,
|
||||
self.norm_b2,
|
||||
self.norm_eps2,
|
||||
)
|
||||
x[i, idx, :] = x_item.squeeze(0)
|
||||
x = self.to_mask(x, padding_mask)
|
||||
else:
|
||||
x = x + attn
|
||||
x = F.layer_norm(x, [self.hidden_dim], self.norm_w1, self.norm_b1, self.norm_eps1)
|
||||
x = x + self.mlp.forward(x)
|
||||
x = F.layer_norm(
|
||||
x,
|
||||
[self.hidden_dim],
|
||||
self.norm_w2,
|
||||
self.norm_b2,
|
||||
self.norm_eps2,
|
||||
)
|
||||
x = x + attn
|
||||
x = F.layer_norm(x, [self.hidden_dim], self.norm_w1, self.norm_b1, self.norm_eps1)
|
||||
x = x + self.mlp.forward(x)
|
||||
x = F.layer_norm(
|
||||
x,
|
||||
[self.hidden_dim],
|
||||
self.norm_w2,
|
||||
self.norm_b2,
|
||||
self.norm_eps2,
|
||||
)
|
||||
return x, k_cache, v_cache
|
||||
|
||||
def decode_next_token(self, x: torch.Tensor, k_cache: torch.Tensor, v_cache: torch.Tensor):
|
||||
@@ -289,8 +294,9 @@ class T2SBlock:
|
||||
|
||||
attn = F.scaled_dot_product_attention(q, k, v)
|
||||
|
||||
attn = attn.permute(2, 0, 1, 3).reshape(batch_size * q_len, self.hidden_dim)
|
||||
attn = attn.view(q_len, batch_size, self.hidden_dim).transpose(1, 0)
|
||||
# attn = attn.permute(2, 0, 1, 3).reshape(batch_size * q_len, self.hidden_dim)
|
||||
# attn = attn.view(q_len, batch_size, self.hidden_dim).transpose(1, 0)
|
||||
attn = attn.transpose(1, 2).reshape(batch_size, q_len, -1)
|
||||
attn = F.linear(attn, self.out_w, self.out_b)
|
||||
|
||||
x = x + attn
|
||||
@@ -328,15 +334,22 @@ class T2STransformer:
|
||||
|
||||
|
||||
class VitsModel(nn.Module):
|
||||
def __init__(self, vits_path):
|
||||
def __init__(self, vits_path, version=None, is_half=True, device="cpu"):
|
||||
super().__init__()
|
||||
# dict_s2 = torch.load(vits_path,map_location="cpu")
|
||||
dict_s2 = torch.load(vits_path, weights_only=False)
|
||||
dict_s2 = load_sovits_new(vits_path)
|
||||
self.hps = dict_s2["config"]
|
||||
if dict_s2["weight"]["enc_p.text_embedding.weight"].shape[0] == 322:
|
||||
self.hps["model"]["version"] = "v1"
|
||||
|
||||
if version is None:
|
||||
if dict_s2["weight"]["enc_p.text_embedding.weight"].shape[0] == 322:
|
||||
self.hps["model"]["version"] = "v1"
|
||||
else:
|
||||
self.hps["model"]["version"] = "v2"
|
||||
else:
|
||||
self.hps["model"]["version"] = "v2"
|
||||
if version in ["v1", "v2", "v3", "v4", "v2Pro", "v2ProPlus"]:
|
||||
self.hps["model"]["version"] = version
|
||||
else:
|
||||
raise ValueError(f"Unsupported version: {version}")
|
||||
|
||||
self.hps = DictToAttrRecursive(self.hps)
|
||||
self.hps.model.semantic_frame_rate = "25hz"
|
||||
@@ -346,11 +359,19 @@ class VitsModel(nn.Module):
|
||||
n_speakers=self.hps.data.n_speakers,
|
||||
**self.hps.model,
|
||||
)
|
||||
self.vq_model.eval()
|
||||
self.vq_model.load_state_dict(dict_s2["weight"], strict=False)
|
||||
self.vq_model.dec.remove_weight_norm()
|
||||
if is_half:
|
||||
self.vq_model = self.vq_model.half()
|
||||
self.vq_model = self.vq_model.to(device)
|
||||
self.vq_model.eval()
|
||||
self.hann_window = torch.hann_window(
|
||||
self.hps.data.win_length, device=device, dtype=torch.float16 if is_half else torch.float32
|
||||
)
|
||||
|
||||
def forward(self, text_seq, pred_semantic, ref_audio, speed=1.0):
|
||||
def forward(self, text_seq, pred_semantic, ref_audio, speed=1.0, sv_emb=None):
|
||||
refer = spectrogram_torch(
|
||||
self.hann_window,
|
||||
ref_audio,
|
||||
self.hps.data.filter_length,
|
||||
self.hps.data.sampling_rate,
|
||||
@@ -358,7 +379,7 @@ class VitsModel(nn.Module):
|
||||
self.hps.data.win_length,
|
||||
center=False,
|
||||
)
|
||||
return self.vq_model(pred_semantic, text_seq, refer, speed)[0, 0]
|
||||
return self.vq_model(pred_semantic, text_seq, refer, speed=speed, sv_emb=sv_emb)[0, 0]
|
||||
|
||||
|
||||
class T2SModel(nn.Module):
|
||||
@@ -433,6 +454,10 @@ class T2SModel(nn.Module):
|
||||
bert = bert.unsqueeze(0)
|
||||
|
||||
x = self.ar_text_embedding(all_phoneme_ids)
|
||||
|
||||
# avoid dtype inconsistency when exporting
|
||||
bert = bert.to(dtype=self.bert_proj.weight.dtype)
|
||||
|
||||
x = x + self.bert_proj(bert.transpose(1, 2))
|
||||
x: torch.Tensor = self.ar_text_position(x)
|
||||
|
||||
@@ -632,7 +657,9 @@ def export(gpt_path, vits_path, ref_audio_path, ref_text, output_path, export_be
|
||||
ref_seq = torch.LongTensor([ref_seq_id]).to(device)
|
||||
ref_bert = ref_bert_T.T.to(ref_seq.device)
|
||||
text_seq_id, text_bert_T, norm_text = get_phones_and_bert(
|
||||
"这是一条测试语音,说什么无所谓,只是给它一个例子", "all_zh", "v2"
|
||||
"这是一个简单的示例,真没想到这么简单就完成了。The King and His Stories.Once there was a king. He likes to write stories, but his stories were not good. As people were afraid of him, they all said his stories were good.After reading them, the writer at once turned to the soldiers and said: Take me back to prison, please.",
|
||||
"auto",
|
||||
"v2",
|
||||
)
|
||||
text_seq = torch.LongTensor([text_seq_id]).to(device)
|
||||
text_bert = text_bert_T.T.to(text_seq.device)
|
||||
@@ -640,7 +667,7 @@ def export(gpt_path, vits_path, ref_audio_path, ref_text, output_path, export_be
|
||||
ssl_content = ssl(ref_audio).to(device)
|
||||
|
||||
# vits_path = "SoVITS_weights_v2/xw_e8_s216.pth"
|
||||
vits = VitsModel(vits_path).to(device)
|
||||
vits = VitsModel(vits_path, device=device, is_half=False)
|
||||
vits.eval()
|
||||
|
||||
# gpt_path = "GPT_weights_v2/xw-e15.ckpt"
|
||||
@@ -679,6 +706,124 @@ def export(gpt_path, vits_path, ref_audio_path, ref_text, output_path, export_be
|
||||
print("#### exported gpt_sovits ####")
|
||||
|
||||
|
||||
def export_prov2(
|
||||
gpt_path,
|
||||
vits_path,
|
||||
version,
|
||||
ref_audio_path,
|
||||
ref_text,
|
||||
output_path,
|
||||
export_bert_and_ssl=False,
|
||||
device="cpu",
|
||||
is_half=True,
|
||||
):
|
||||
if sv_cn_model == None:
|
||||
init_sv_cn(device, is_half)
|
||||
|
||||
if not os.path.exists(output_path):
|
||||
os.makedirs(output_path)
|
||||
print(f"目录已创建: {output_path}")
|
||||
else:
|
||||
print(f"目录已存在: {output_path}")
|
||||
|
||||
ref_audio = torch.tensor([load_audio(ref_audio_path, 16000)]).float()
|
||||
ssl = SSLModel()
|
||||
if export_bert_and_ssl:
|
||||
s = ExportSSLModel(torch.jit.trace(ssl, example_inputs=(ref_audio)))
|
||||
ssl_path = os.path.join(output_path, "ssl_model.pt")
|
||||
torch.jit.script(s).save(ssl_path)
|
||||
print("#### exported ssl ####")
|
||||
export_bert(output_path)
|
||||
else:
|
||||
s = ExportSSLModel(ssl)
|
||||
|
||||
print(f"device: {device}")
|
||||
|
||||
ref_seq_id, ref_bert_T, ref_norm_text = get_phones_and_bert(ref_text, "all_zh", "v2")
|
||||
ref_seq = torch.LongTensor([ref_seq_id]).to(device)
|
||||
ref_bert = ref_bert_T.T
|
||||
if is_half:
|
||||
ref_bert = ref_bert.half()
|
||||
ref_bert = ref_bert.to(ref_seq.device)
|
||||
|
||||
text_seq_id, text_bert_T, norm_text = get_phones_and_bert(
|
||||
"这是一个简单的示例,真没想到这么简单就完成了。The King and His Stories.Once there was a king. He likes to write stories, but his stories were not good. As people were afraid of him, they all said his stories were good.After reading them, the writer at once turned to the soldiers and said: Take me back to prison, please.",
|
||||
"auto",
|
||||
"v2",
|
||||
)
|
||||
text_seq = torch.LongTensor([text_seq_id]).to(device)
|
||||
text_bert = text_bert_T.T
|
||||
if is_half:
|
||||
text_bert = text_bert.half()
|
||||
text_bert = text_bert.to(text_seq.device)
|
||||
|
||||
ssl_content = ssl(ref_audio)
|
||||
if is_half:
|
||||
ssl_content = ssl_content.half()
|
||||
ssl_content = ssl_content.to(device)
|
||||
|
||||
sv_model = ExportERes2NetV2(sv_cn_model)
|
||||
|
||||
# vits_path = "SoVITS_weights_v2/xw_e8_s216.pth"
|
||||
vits = VitsModel(vits_path, version, is_half=is_half, device=device)
|
||||
vits.eval()
|
||||
|
||||
# gpt_path = "GPT_weights_v2/xw-e15.ckpt"
|
||||
# dict_s1 = torch.load(gpt_path, map_location=device)
|
||||
dict_s1 = torch.load(gpt_path, weights_only=False)
|
||||
raw_t2s = get_raw_t2s_model(dict_s1).to(device)
|
||||
print("#### get_raw_t2s_model ####")
|
||||
print(raw_t2s.config)
|
||||
if is_half:
|
||||
raw_t2s = raw_t2s.half()
|
||||
t2s_m = T2SModel(raw_t2s)
|
||||
t2s_m.eval()
|
||||
t2s = torch.jit.script(t2s_m).to(device)
|
||||
print("#### script t2s_m ####")
|
||||
|
||||
print("vits.hps.data.sampling_rate:", vits.hps.data.sampling_rate)
|
||||
gpt_sovits = GPT_SoVITS_V2Pro(t2s, vits, sv_model).to(device)
|
||||
gpt_sovits.eval()
|
||||
|
||||
ref_audio_sr = s.resample(ref_audio, 16000, 32000)
|
||||
if is_half:
|
||||
ref_audio_sr = ref_audio_sr.half()
|
||||
ref_audio_sr = ref_audio_sr.to(device)
|
||||
|
||||
torch._dynamo.mark_dynamic(ssl_content, 2)
|
||||
torch._dynamo.mark_dynamic(ref_audio_sr, 1)
|
||||
torch._dynamo.mark_dynamic(ref_seq, 1)
|
||||
torch._dynamo.mark_dynamic(text_seq, 1)
|
||||
torch._dynamo.mark_dynamic(ref_bert, 0)
|
||||
torch._dynamo.mark_dynamic(text_bert, 0)
|
||||
# torch._dynamo.mark_dynamic(sv_emb, 0)
|
||||
|
||||
top_k = torch.LongTensor([5]).to(device)
|
||||
# 先跑一遍 sv_model 让它加载 cache,详情见 L880
|
||||
gpt_sovits.sv_model(ref_audio_sr)
|
||||
|
||||
with torch.no_grad():
|
||||
gpt_sovits_export = torch.jit.trace(
|
||||
gpt_sovits,
|
||||
example_inputs=(
|
||||
ssl_content,
|
||||
ref_audio_sr,
|
||||
ref_seq,
|
||||
text_seq,
|
||||
ref_bert,
|
||||
text_bert,
|
||||
top_k,
|
||||
),
|
||||
)
|
||||
|
||||
gpt_sovits_path = os.path.join(output_path, "gpt_sovits_model.pt")
|
||||
gpt_sovits_export.save(gpt_sovits_path)
|
||||
print("#### exported gpt_sovits ####")
|
||||
audio = gpt_sovits_export(ssl_content, ref_audio_sr, ref_seq, text_seq, ref_bert, text_bert, top_k)
|
||||
print("start write wav")
|
||||
soundfile.write("out.wav", audio.float().detach().cpu().numpy(), 32000)
|
||||
|
||||
|
||||
@torch.jit.script
|
||||
def parse_audio(ref_audio):
|
||||
ref_audio_16k = torchaudio.functional.resample(ref_audio, 48000, 16000).float() # .to(ref_audio.device)
|
||||
@@ -717,6 +862,67 @@ class GPT_SoVITS(nn.Module):
|
||||
return audio
|
||||
|
||||
|
||||
class ExportERes2NetV2(nn.Module):
|
||||
def __init__(self, sv_cn_model: SV):
|
||||
super(ExportERes2NetV2, self).__init__()
|
||||
self.bn1 = sv_cn_model.embedding_model.bn1
|
||||
self.conv1 = sv_cn_model.embedding_model.conv1
|
||||
self.layer1 = sv_cn_model.embedding_model.layer1
|
||||
self.layer2 = sv_cn_model.embedding_model.layer2
|
||||
self.layer3 = sv_cn_model.embedding_model.layer3
|
||||
self.layer4 = sv_cn_model.embedding_model.layer4
|
||||
self.layer3_ds = sv_cn_model.embedding_model.layer3_ds
|
||||
self.fuse34 = sv_cn_model.embedding_model.fuse34
|
||||
|
||||
# audio_16k.shape: [1,N]
|
||||
def forward(self, audio_16k):
|
||||
# 这个 fbank 函数有一个 cache, 不过不要紧,它跟 audio_16k 的长度无关
|
||||
# 只跟 device 和 dtype 有关
|
||||
x = Kaldi.fbank(audio_16k, num_mel_bins=80, sample_frequency=16000, dither=0)
|
||||
x = torch.stack([x])
|
||||
|
||||
x = x.permute(0, 2, 1) # (B,T,F) => (B,F,T)
|
||||
x = x.unsqueeze_(1)
|
||||
out = F.relu(self.bn1(self.conv1(x)))
|
||||
out1 = self.layer1(out)
|
||||
out2 = self.layer2(out1)
|
||||
out3 = self.layer3(out2)
|
||||
out4 = self.layer4(out3)
|
||||
out3_ds = self.layer3_ds(out3)
|
||||
fuse_out34 = self.fuse34(out4, out3_ds)
|
||||
return fuse_out34.flatten(start_dim=1, end_dim=2).mean(-1)
|
||||
|
||||
|
||||
class GPT_SoVITS_V2Pro(nn.Module):
|
||||
def __init__(self, t2s: T2SModel, vits: VitsModel, sv_model: ExportERes2NetV2):
|
||||
super().__init__()
|
||||
self.t2s = t2s
|
||||
self.vits = vits
|
||||
self.sv_model = sv_model
|
||||
|
||||
def forward(
|
||||
self,
|
||||
ssl_content: torch.Tensor,
|
||||
ref_audio_sr: torch.Tensor,
|
||||
ref_seq: Tensor,
|
||||
text_seq: Tensor,
|
||||
ref_bert: Tensor,
|
||||
text_bert: Tensor,
|
||||
top_k: LongTensor,
|
||||
speed=1.0,
|
||||
):
|
||||
codes = self.vits.vq_model.extract_latent(ssl_content)
|
||||
prompt_semantic = codes[0, 0]
|
||||
prompts = prompt_semantic.unsqueeze(0)
|
||||
|
||||
audio_16k = resamplex(ref_audio_sr, 32000, 16000).to(ref_audio_sr.dtype)
|
||||
sv_emb = self.sv_model(audio_16k)
|
||||
|
||||
pred_semantic = self.t2s(prompts, ref_seq, text_seq, ref_bert, text_bert, top_k)
|
||||
audio = self.vits(text_seq, pred_semantic, ref_audio_sr, speed, sv_emb)
|
||||
return audio
|
||||
|
||||
|
||||
def test():
|
||||
parser = argparse.ArgumentParser(description="GPT-SoVITS Command Line Tool")
|
||||
parser.add_argument("--gpt_model", required=True, help="Path to the GPT model file")
|
||||
@@ -839,23 +1045,37 @@ def main():
|
||||
parser.add_argument("--output_path", required=True, help="Path to the output directory")
|
||||
parser.add_argument("--export_common_model", action="store_true", help="Export Bert and SSL model")
|
||||
parser.add_argument("--device", help="Device to use")
|
||||
parser.add_argument("--version", help="version of the model", default="v2")
|
||||
parser.add_argument("--no-half", action="store_true", help="Do not use half precision for model weights")
|
||||
|
||||
args = parser.parse_args()
|
||||
export(
|
||||
gpt_path=args.gpt_model,
|
||||
vits_path=args.sovits_model,
|
||||
ref_audio_path=args.ref_audio,
|
||||
ref_text=args.ref_text,
|
||||
output_path=args.output_path,
|
||||
device=args.device,
|
||||
export_bert_and_ssl=args.export_common_model,
|
||||
)
|
||||
if args.version in ["v2Pro", "v2ProPlus"]:
|
||||
is_half = not args.no_half
|
||||
print(f"Using half precision: {is_half}")
|
||||
export_prov2(
|
||||
gpt_path=args.gpt_model,
|
||||
vits_path=args.sovits_model,
|
||||
version=args.version,
|
||||
ref_audio_path=args.ref_audio,
|
||||
ref_text=args.ref_text,
|
||||
output_path=args.output_path,
|
||||
export_bert_and_ssl=args.export_common_model,
|
||||
device=args.device,
|
||||
is_half=is_half,
|
||||
)
|
||||
else:
|
||||
export(
|
||||
gpt_path=args.gpt_model,
|
||||
vits_path=args.sovits_model,
|
||||
ref_audio_path=args.ref_audio,
|
||||
ref_text=args.ref_text,
|
||||
output_path=args.output_path,
|
||||
device=args.device,
|
||||
export_bert_and_ssl=args.export_common_model,
|
||||
)
|
||||
|
||||
|
||||
import inference_webui
|
||||
|
||||
if __name__ == "__main__":
|
||||
inference_webui.is_half = False
|
||||
inference_webui.dtype = torch.float32
|
||||
main()
|
||||
with torch.no_grad():
|
||||
main()
|
||||
# test()
|
||||
|
||||
@@ -243,6 +243,7 @@ class ExportGPTSovitsHalf(torch.nn.Module):
|
||||
self.sampling_rate: int = hps.data.sampling_rate
|
||||
self.hop_length: int = hps.data.hop_length
|
||||
self.win_length: int = hps.data.win_length
|
||||
self.hann_window = torch.hann_window(self.win_length, device=device, dtype=torch.float32)
|
||||
|
||||
def forward(
|
||||
self,
|
||||
@@ -255,6 +256,7 @@ class ExportGPTSovitsHalf(torch.nn.Module):
|
||||
top_k,
|
||||
):
|
||||
refer = spectrogram_torch(
|
||||
self.hann_window,
|
||||
ref_audio_32k,
|
||||
self.filter_length,
|
||||
self.sampling_rate,
|
||||
@@ -321,6 +323,7 @@ class ExportGPTSovitsV4Half(torch.nn.Module):
|
||||
self.sampling_rate: int = hps.data.sampling_rate
|
||||
self.hop_length: int = hps.data.hop_length
|
||||
self.win_length: int = hps.data.win_length
|
||||
self.hann_window = torch.hann_window(self.win_length, device=device, dtype=torch.float32)
|
||||
|
||||
def forward(
|
||||
self,
|
||||
@@ -333,6 +336,7 @@ class ExportGPTSovitsV4Half(torch.nn.Module):
|
||||
top_k,
|
||||
):
|
||||
refer = spectrogram_torch(
|
||||
self.hann_window,
|
||||
ref_audio_32k,
|
||||
self.filter_length,
|
||||
self.sampling_rate,
|
||||
@@ -402,7 +406,7 @@ class GPTSoVITSV3(torch.nn.Module):
|
||||
chunk_len = 934 - fea_ref.shape[2]
|
||||
wav_gen_list = []
|
||||
idx = 0
|
||||
fea_todo = fea_todo[:,:,:-5]
|
||||
fea_todo = fea_todo[:, :, :-5]
|
||||
wav_gen_length = fea_todo.shape[2] * 256
|
||||
while 1:
|
||||
# current_time = datetime.now()
|
||||
@@ -434,7 +438,8 @@ class GPTSoVITSV3(torch.nn.Module):
|
||||
|
||||
wav_gen = torch.cat(wav_gen_list, 2)
|
||||
return wav_gen[0][0][:wav_gen_length]
|
||||
|
||||
|
||||
|
||||
class GPTSoVITSV4(torch.nn.Module):
|
||||
def __init__(self, gpt_sovits_half, cfm, hifigan):
|
||||
super().__init__()
|
||||
@@ -461,7 +466,7 @@ class GPTSoVITSV4(torch.nn.Module):
|
||||
chunk_len = 1000 - fea_ref.shape[2]
|
||||
wav_gen_list = []
|
||||
idx = 0
|
||||
fea_todo = fea_todo[:,:,:-10]
|
||||
fea_todo = fea_todo[:, :, :-10]
|
||||
wav_gen_length = fea_todo.shape[2] * 480
|
||||
while 1:
|
||||
# current_time = datetime.now()
|
||||
@@ -501,8 +506,7 @@ def init_bigvgan():
|
||||
|
||||
bigvgan_model = bigvgan.BigVGAN.from_pretrained(
|
||||
"%s/GPT_SoVITS/pretrained_models/models--nvidia--bigvgan_v2_24khz_100band_256x" % (now_dir,),
|
||||
use_cuda_kernel=False,
|
||||
) # if True, RuntimeError: Ninja is required to load C++ extensions
|
||||
)
|
||||
# remove weight norm in the model and set to eval mode
|
||||
bigvgan_model.remove_weight_norm()
|
||||
bigvgan_model = bigvgan_model.eval()
|
||||
@@ -577,6 +581,7 @@ from process_ckpt import get_sovits_version_from_path_fast, load_sovits_new
|
||||
|
||||
v3v4set = {"v3", "v4"}
|
||||
|
||||
|
||||
def get_sovits_weights(sovits_path):
|
||||
path_sovits_v3 = "GPT_SoVITS/pretrained_models/s2Gv3.pth"
|
||||
is_exist_s2gv3 = os.path.exists(path_sovits_v3)
|
||||
@@ -699,14 +704,13 @@ def export_cfm(
|
||||
return export_cfm
|
||||
|
||||
|
||||
def export_1(ref_wav_path,ref_wav_text,version="v3"):
|
||||
def export_1(ref_wav_path, ref_wav_text, version="v3"):
|
||||
if version == "v3":
|
||||
sovits = get_sovits_weights("GPT_SoVITS/pretrained_models/s2Gv3.pth")
|
||||
init_bigvgan()
|
||||
else:
|
||||
sovits = get_sovits_weights("GPT_SoVITS/pretrained_models/gsv-v4-pretrained/s2Gv4.pth")
|
||||
init_hifigan()
|
||||
|
||||
|
||||
dict_s1 = torch.load("GPT_SoVITS/pretrained_models/s1v3.ckpt")
|
||||
raw_t2s = get_raw_t2s_model(dict_s1).to(device)
|
||||
@@ -751,9 +755,7 @@ def export_1(ref_wav_path,ref_wav_text,version="v3"):
|
||||
# phones1, bert1, norm_text1 = get_phones_and_bert(
|
||||
# "你这老坏蛋,我找了你这么久,真没想到在这里找到你。他说。", "all_zh", "v3"
|
||||
# )
|
||||
phones1, bert1, norm_text1 = get_phones_and_bert(
|
||||
ref_wav_text, "auto", "v3"
|
||||
)
|
||||
phones1, bert1, norm_text1 = get_phones_and_bert(ref_wav_text, "auto", "v3")
|
||||
phones2, bert2, norm_text2 = get_phones_and_bert(
|
||||
"这是一个简单的示例,真没想到这么简单就完成了。The King and His Stories.Once there was a king. He likes to write stories, but his stories were not good. As people were afraid of him, they all said his stories were good.After reading them, the writer at once turned to the soldiers and said: Take me back to prison, please.",
|
||||
"auto",
|
||||
@@ -914,7 +916,7 @@ def export_1(ref_wav_path,ref_wav_text,version="v3"):
|
||||
hifigan_model_ = torch.jit.trace(hifigan_model, optimize=True, example_inputs=(cmf_res_rand,))
|
||||
hifigan_model_.save("onnx/ad/hifigan_model.pt")
|
||||
wav_gen = hifigan_model(cmf_res)
|
||||
|
||||
|
||||
print("wav_gen:", wav_gen.shape, wav_gen.dtype)
|
||||
audio = wav_gen[0][0].cpu().detach().numpy()
|
||||
|
||||
@@ -1149,7 +1151,7 @@ def export_2(version="v3"):
|
||||
raw_t2s = raw_t2s.half().to(device)
|
||||
t2s_m = T2SModel(raw_t2s).half().to(device)
|
||||
t2s_m.eval()
|
||||
t2s_m = torch.jit.script(t2s_m)
|
||||
t2s_m = torch.jit.script(t2s_m).to(device)
|
||||
t2s_m.eval()
|
||||
# t2s_m.top_k = 15
|
||||
logger.info("t2s_m ok")
|
||||
@@ -1201,7 +1203,6 @@ def export_2(version="v3"):
|
||||
gpt_sovits_v3v4 = gpt_sovits_v3 if version == "v3" else gpt_sovits_v4
|
||||
sr = 24000 if version == "v3" else 48000
|
||||
|
||||
|
||||
time.sleep(5)
|
||||
# print("thread:", torch.get_num_threads())
|
||||
# print("thread:", torch.get_num_interop_threads())
|
||||
@@ -1212,14 +1213,14 @@ def export_2(version="v3"):
|
||||
"汗流浃背了呀!老弟~ My uncle has two dogs. One is big and the other is small. He likes them very much. He often plays with them. He takes them for a walk every day. He says they are his good friends. He is very happy with them. 最后还是我得了 MVP....",
|
||||
gpt_sovits_v3v4,
|
||||
"out.wav",
|
||||
sr
|
||||
sr,
|
||||
)
|
||||
|
||||
test_export(
|
||||
"你小子是什么来路.汗流浃背了呀!老弟~ My uncle has two dogs. He is very happy with them. 最后还是我得了 MVP!",
|
||||
gpt_sovits_v3v4,
|
||||
"out2.wav",
|
||||
sr
|
||||
sr,
|
||||
)
|
||||
|
||||
# test_export(
|
||||
@@ -1251,6 +1252,6 @@ def test_export_gpt_sovits_v3():
|
||||
|
||||
|
||||
with torch.no_grad():
|
||||
export_1("onnx/ad/ref.wav","你这老坏蛋,我找了你这么久,真没想到在这里找到你。他说。","v4")
|
||||
# export_2("v4")
|
||||
# export_1("onnx/ad/ref.wav","你这老坏蛋,我找了你这么久,真没想到在这里找到你。他说。","v4")
|
||||
export_2("v4")
|
||||
# test_export_gpt_sovits_v3()
|
||||
|
||||
@@ -1,13 +1 @@
|
||||
# from f5_tts.model.cfm import CFM
|
||||
#
|
||||
# from f5_tts.model.backbones.unett import UNetT
|
||||
from GPT_SoVITS.f5_tts.model.backbones.dit import DiT
|
||||
# from f5_tts.model.backbones.dit import DiTNoCond
|
||||
# from f5_tts.model.backbones.dit import DiTNoCondNoT
|
||||
# from f5_tts.model.backbones.mmdit import MMDiT
|
||||
|
||||
# from f5_tts.model.trainer import Trainer
|
||||
|
||||
|
||||
# __all__ = ["CFM", "UNetT", "DiT", "MMDiT", "Trainer"]
|
||||
# __all__ = ["CFM", "UNetT", "DiTNoCond","DiT", "MMDiT"]
|
||||
|
||||
@@ -143,9 +143,9 @@ class DiT(nn.Module):
|
||||
drop_audio_cond=False, # cfg for cond audio
|
||||
drop_text=False, # cfg for text
|
||||
# mask: bool["b n"] | None = None, # noqa: F722
|
||||
infer=False, # bool
|
||||
text_cache=None, # torch tensor as text_embed
|
||||
dt_cache=None, # torch tensor as dt
|
||||
infer=False, # bool
|
||||
text_cache=None, # torch tensor as text_embed
|
||||
dt_cache=None, # torch tensor as dt
|
||||
):
|
||||
x = x0.transpose(2, 1)
|
||||
cond = cond0.transpose(2, 1)
|
||||
@@ -191,4 +191,4 @@ class DiT(nn.Module):
|
||||
if infer:
|
||||
return output, text_embed, dt
|
||||
else:
|
||||
return output
|
||||
return output
|
||||
|
||||
@@ -1,146 +0,0 @@
|
||||
"""
|
||||
ein notation:
|
||||
b - batch
|
||||
n - sequence
|
||||
nt - text sequence
|
||||
nw - raw wave length
|
||||
d - dimension
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import torch
|
||||
from torch import nn
|
||||
|
||||
from x_transformers.x_transformers import RotaryEmbedding
|
||||
|
||||
from f5_tts.model.modules import (
|
||||
TimestepEmbedding,
|
||||
ConvPositionEmbedding,
|
||||
MMDiTBlock,
|
||||
AdaLayerNormZero_Final,
|
||||
precompute_freqs_cis,
|
||||
get_pos_embed_indices,
|
||||
)
|
||||
|
||||
|
||||
# text embedding
|
||||
|
||||
|
||||
class TextEmbedding(nn.Module):
|
||||
def __init__(self, out_dim, text_num_embeds):
|
||||
super().__init__()
|
||||
self.text_embed = nn.Embedding(text_num_embeds + 1, out_dim) # will use 0 as filler token
|
||||
|
||||
self.precompute_max_pos = 1024
|
||||
self.register_buffer("freqs_cis", precompute_freqs_cis(out_dim, self.precompute_max_pos), persistent=False)
|
||||
|
||||
def forward(self, text: int["b nt"], drop_text=False) -> int["b nt d"]: # noqa: F722
|
||||
text = text + 1
|
||||
if drop_text:
|
||||
text = torch.zeros_like(text)
|
||||
text = self.text_embed(text)
|
||||
|
||||
# sinus pos emb
|
||||
batch_start = torch.zeros((text.shape[0],), dtype=torch.long)
|
||||
batch_text_len = text.shape[1]
|
||||
pos_idx = get_pos_embed_indices(batch_start, batch_text_len, max_pos=self.precompute_max_pos)
|
||||
text_pos_embed = self.freqs_cis[pos_idx]
|
||||
|
||||
text = text + text_pos_embed
|
||||
|
||||
return text
|
||||
|
||||
|
||||
# noised input & masked cond audio embedding
|
||||
|
||||
|
||||
class AudioEmbedding(nn.Module):
|
||||
def __init__(self, in_dim, out_dim):
|
||||
super().__init__()
|
||||
self.linear = nn.Linear(2 * in_dim, out_dim)
|
||||
self.conv_pos_embed = ConvPositionEmbedding(out_dim)
|
||||
|
||||
def forward(self, x: float["b n d"], cond: float["b n d"], drop_audio_cond=False): # noqa: F722
|
||||
if drop_audio_cond:
|
||||
cond = torch.zeros_like(cond)
|
||||
x = torch.cat((x, cond), dim=-1)
|
||||
x = self.linear(x)
|
||||
x = self.conv_pos_embed(x) + x
|
||||
return x
|
||||
|
||||
|
||||
# Transformer backbone using MM-DiT blocks
|
||||
|
||||
|
||||
class MMDiT(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
*,
|
||||
dim,
|
||||
depth=8,
|
||||
heads=8,
|
||||
dim_head=64,
|
||||
dropout=0.1,
|
||||
ff_mult=4,
|
||||
text_num_embeds=256,
|
||||
mel_dim=100,
|
||||
):
|
||||
super().__init__()
|
||||
|
||||
self.time_embed = TimestepEmbedding(dim)
|
||||
self.text_embed = TextEmbedding(dim, text_num_embeds)
|
||||
self.audio_embed = AudioEmbedding(mel_dim, dim)
|
||||
|
||||
self.rotary_embed = RotaryEmbedding(dim_head)
|
||||
|
||||
self.dim = dim
|
||||
self.depth = depth
|
||||
|
||||
self.transformer_blocks = nn.ModuleList(
|
||||
[
|
||||
MMDiTBlock(
|
||||
dim=dim,
|
||||
heads=heads,
|
||||
dim_head=dim_head,
|
||||
dropout=dropout,
|
||||
ff_mult=ff_mult,
|
||||
context_pre_only=i == depth - 1,
|
||||
)
|
||||
for i in range(depth)
|
||||
]
|
||||
)
|
||||
self.norm_out = AdaLayerNormZero_Final(dim) # final modulation
|
||||
self.proj_out = nn.Linear(dim, mel_dim)
|
||||
|
||||
def forward(
|
||||
self,
|
||||
x: float["b n d"], # nosied input audio # noqa: F722
|
||||
cond: float["b n d"], # masked cond audio # noqa: F722
|
||||
text: int["b nt"], # text # noqa: F722
|
||||
time: float["b"] | float[""], # time step # noqa: F821 F722
|
||||
drop_audio_cond, # cfg for cond audio
|
||||
drop_text, # cfg for text
|
||||
mask: bool["b n"] | None = None, # noqa: F722
|
||||
):
|
||||
batch = x.shape[0]
|
||||
if time.ndim == 0:
|
||||
time = time.repeat(batch)
|
||||
|
||||
# t: conditioning (time), c: context (text + masked cond audio), x: noised input audio
|
||||
t = self.time_embed(time)
|
||||
c = self.text_embed(text, drop_text=drop_text)
|
||||
x = self.audio_embed(x, cond, drop_audio_cond=drop_audio_cond)
|
||||
|
||||
seq_len = x.shape[1]
|
||||
text_len = text.shape[1]
|
||||
rope_audio = self.rotary_embed.forward_from_seq_len(seq_len)
|
||||
rope_text = self.rotary_embed.forward_from_seq_len(text_len)
|
||||
|
||||
for block in self.transformer_blocks:
|
||||
c, x = block(x, c, t, mask=mask, rope=rope_audio, c_rope=rope_text)
|
||||
|
||||
x = self.norm_out(x, t)
|
||||
output = self.proj_out(x)
|
||||
|
||||
return output
|
||||
@@ -1,219 +0,0 @@
|
||||
"""
|
||||
ein notation:
|
||||
b - batch
|
||||
n - sequence
|
||||
nt - text sequence
|
||||
nw - raw wave length
|
||||
d - dimension
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
from typing import Literal
|
||||
|
||||
import torch
|
||||
from torch import nn
|
||||
import torch.nn.functional as F
|
||||
|
||||
from x_transformers import RMSNorm
|
||||
from x_transformers.x_transformers import RotaryEmbedding
|
||||
|
||||
from f5_tts.model.modules import (
|
||||
TimestepEmbedding,
|
||||
ConvNeXtV2Block,
|
||||
ConvPositionEmbedding,
|
||||
Attention,
|
||||
AttnProcessor,
|
||||
FeedForward,
|
||||
precompute_freqs_cis,
|
||||
get_pos_embed_indices,
|
||||
)
|
||||
|
||||
|
||||
# Text embedding
|
||||
|
||||
|
||||
class TextEmbedding(nn.Module):
|
||||
def __init__(self, text_num_embeds, text_dim, conv_layers=0, conv_mult=2):
|
||||
super().__init__()
|
||||
self.text_embed = nn.Embedding(text_num_embeds + 1, text_dim) # use 0 as filler token
|
||||
|
||||
if conv_layers > 0:
|
||||
self.extra_modeling = True
|
||||
self.precompute_max_pos = 4096 # ~44s of 24khz audio
|
||||
self.register_buffer("freqs_cis", precompute_freqs_cis(text_dim, self.precompute_max_pos), persistent=False)
|
||||
self.text_blocks = nn.Sequential(
|
||||
*[ConvNeXtV2Block(text_dim, text_dim * conv_mult) for _ in range(conv_layers)]
|
||||
)
|
||||
else:
|
||||
self.extra_modeling = False
|
||||
|
||||
def forward(self, text: int["b nt"], seq_len, drop_text=False): # noqa: F722
|
||||
text = text + 1 # use 0 as filler token. preprocess of batch pad -1, see list_str_to_idx()
|
||||
text = text[:, :seq_len] # curtail if character tokens are more than the mel spec tokens
|
||||
batch, text_len = text.shape[0], text.shape[1]
|
||||
text = F.pad(text, (0, seq_len - text_len), value=0)
|
||||
|
||||
if drop_text: # cfg for text
|
||||
text = torch.zeros_like(text)
|
||||
|
||||
text = self.text_embed(text) # b n -> b n d
|
||||
|
||||
# possible extra modeling
|
||||
if self.extra_modeling:
|
||||
# sinus pos emb
|
||||
batch_start = torch.zeros((batch,), dtype=torch.long)
|
||||
pos_idx = get_pos_embed_indices(batch_start, seq_len, max_pos=self.precompute_max_pos)
|
||||
text_pos_embed = self.freqs_cis[pos_idx]
|
||||
text = text + text_pos_embed
|
||||
|
||||
# convnextv2 blocks
|
||||
text = self.text_blocks(text)
|
||||
|
||||
return text
|
||||
|
||||
|
||||
# noised input audio and context mixing embedding
|
||||
|
||||
|
||||
class InputEmbedding(nn.Module):
|
||||
def __init__(self, mel_dim, text_dim, out_dim):
|
||||
super().__init__()
|
||||
self.proj = nn.Linear(mel_dim * 2 + text_dim, out_dim)
|
||||
self.conv_pos_embed = ConvPositionEmbedding(dim=out_dim)
|
||||
|
||||
def forward(self, x: float["b n d"], cond: float["b n d"], text_embed: float["b n d"], drop_audio_cond=False): # noqa: F722
|
||||
if drop_audio_cond: # cfg for cond audio
|
||||
cond = torch.zeros_like(cond)
|
||||
|
||||
x = self.proj(torch.cat((x, cond, text_embed), dim=-1))
|
||||
x = self.conv_pos_embed(x) + x
|
||||
return x
|
||||
|
||||
|
||||
# Flat UNet Transformer backbone
|
||||
|
||||
|
||||
class UNetT(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
*,
|
||||
dim,
|
||||
depth=8,
|
||||
heads=8,
|
||||
dim_head=64,
|
||||
dropout=0.1,
|
||||
ff_mult=4,
|
||||
mel_dim=100,
|
||||
text_num_embeds=256,
|
||||
text_dim=None,
|
||||
conv_layers=0,
|
||||
skip_connect_type: Literal["add", "concat", "none"] = "concat",
|
||||
):
|
||||
super().__init__()
|
||||
assert depth % 2 == 0, "UNet-Transformer's depth should be even."
|
||||
|
||||
self.time_embed = TimestepEmbedding(dim)
|
||||
if text_dim is None:
|
||||
text_dim = mel_dim
|
||||
self.text_embed = TextEmbedding(text_num_embeds, text_dim, conv_layers=conv_layers)
|
||||
self.input_embed = InputEmbedding(mel_dim, text_dim, dim)
|
||||
|
||||
self.rotary_embed = RotaryEmbedding(dim_head)
|
||||
|
||||
# transformer layers & skip connections
|
||||
|
||||
self.dim = dim
|
||||
self.skip_connect_type = skip_connect_type
|
||||
needs_skip_proj = skip_connect_type == "concat"
|
||||
|
||||
self.depth = depth
|
||||
self.layers = nn.ModuleList([])
|
||||
|
||||
for idx in range(depth):
|
||||
is_later_half = idx >= (depth // 2)
|
||||
|
||||
attn_norm = RMSNorm(dim)
|
||||
attn = Attention(
|
||||
processor=AttnProcessor(),
|
||||
dim=dim,
|
||||
heads=heads,
|
||||
dim_head=dim_head,
|
||||
dropout=dropout,
|
||||
)
|
||||
|
||||
ff_norm = RMSNorm(dim)
|
||||
ff = FeedForward(dim=dim, mult=ff_mult, dropout=dropout, approximate="tanh")
|
||||
|
||||
skip_proj = nn.Linear(dim * 2, dim, bias=False) if needs_skip_proj and is_later_half else None
|
||||
|
||||
self.layers.append(
|
||||
nn.ModuleList(
|
||||
[
|
||||
skip_proj,
|
||||
attn_norm,
|
||||
attn,
|
||||
ff_norm,
|
||||
ff,
|
||||
]
|
||||
)
|
||||
)
|
||||
|
||||
self.norm_out = RMSNorm(dim)
|
||||
self.proj_out = nn.Linear(dim, mel_dim)
|
||||
|
||||
def forward(
|
||||
self,
|
||||
x: float["b n d"], # nosied input audio # noqa: F722
|
||||
cond: float["b n d"], # masked cond audio # noqa: F722
|
||||
text: int["b nt"], # text # noqa: F722
|
||||
time: float["b"] | float[""], # time step # noqa: F821 F722
|
||||
drop_audio_cond, # cfg for cond audio
|
||||
drop_text, # cfg for text
|
||||
mask: bool["b n"] | None = None, # noqa: F722
|
||||
):
|
||||
batch, seq_len = x.shape[0], x.shape[1]
|
||||
if time.ndim == 0:
|
||||
time = time.repeat(batch)
|
||||
|
||||
# t: conditioning time, c: context (text + masked cond audio), x: noised input audio
|
||||
t = self.time_embed(time)
|
||||
text_embed = self.text_embed(text, seq_len, drop_text=drop_text)
|
||||
x = self.input_embed(x, cond, text_embed, drop_audio_cond=drop_audio_cond)
|
||||
|
||||
# postfix time t to input x, [b n d] -> [b n+1 d]
|
||||
x = torch.cat([t.unsqueeze(1), x], dim=1) # pack t to x
|
||||
if mask is not None:
|
||||
mask = F.pad(mask, (1, 0), value=1)
|
||||
|
||||
rope = self.rotary_embed.forward_from_seq_len(seq_len + 1)
|
||||
|
||||
# flat unet transformer
|
||||
skip_connect_type = self.skip_connect_type
|
||||
skips = []
|
||||
for idx, (maybe_skip_proj, attn_norm, attn, ff_norm, ff) in enumerate(self.layers):
|
||||
layer = idx + 1
|
||||
|
||||
# skip connection logic
|
||||
is_first_half = layer <= (self.depth // 2)
|
||||
is_later_half = not is_first_half
|
||||
|
||||
if is_first_half:
|
||||
skips.append(x)
|
||||
|
||||
if is_later_half:
|
||||
skip = skips.pop()
|
||||
if skip_connect_type == "concat":
|
||||
x = torch.cat((x, skip), dim=-1)
|
||||
x = maybe_skip_proj(x)
|
||||
elif skip_connect_type == "add":
|
||||
x = x + skip
|
||||
|
||||
# attention and feedforward blocks
|
||||
x = attn(attn_norm(x), rope=rope, mask=mask) + x
|
||||
x = ff(ff_norm(x)) + x
|
||||
|
||||
assert len(skips) == 0
|
||||
|
||||
x = self.norm_out(x)[:, 1:, :] # unpack t from x
|
||||
|
||||
return self.proj_out(x)
|
||||
@@ -10,139 +10,13 @@ d - dimension
|
||||
from __future__ import annotations
|
||||
|
||||
import math
|
||||
from typing import Optional
|
||||
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
import torchaudio
|
||||
from librosa.filters import mel as librosa_mel_fn
|
||||
from torch import nn
|
||||
from x_transformers.x_transformers import apply_rotary_pos_emb
|
||||
|
||||
|
||||
# raw wav to mel spec
|
||||
|
||||
|
||||
mel_basis_cache = {}
|
||||
hann_window_cache = {}
|
||||
|
||||
|
||||
def get_bigvgan_mel_spectrogram(
|
||||
waveform,
|
||||
n_fft=1024,
|
||||
n_mel_channels=100,
|
||||
target_sample_rate=24000,
|
||||
hop_length=256,
|
||||
win_length=1024,
|
||||
fmin=0,
|
||||
fmax=None,
|
||||
center=False,
|
||||
): # Copy from https://github.com/NVIDIA/BigVGAN/tree/main
|
||||
device = waveform.device
|
||||
key = f"{n_fft}_{n_mel_channels}_{target_sample_rate}_{hop_length}_{win_length}_{fmin}_{fmax}_{device}"
|
||||
|
||||
if key not in mel_basis_cache:
|
||||
mel = librosa_mel_fn(sr=target_sample_rate, n_fft=n_fft, n_mels=n_mel_channels, fmin=fmin, fmax=fmax)
|
||||
mel_basis_cache[key] = torch.from_numpy(mel).float().to(device) # TODO: why they need .float()?
|
||||
hann_window_cache[key] = torch.hann_window(win_length).to(device)
|
||||
|
||||
mel_basis = mel_basis_cache[key]
|
||||
hann_window = hann_window_cache[key]
|
||||
|
||||
padding = (n_fft - hop_length) // 2
|
||||
waveform = torch.nn.functional.pad(waveform.unsqueeze(1), (padding, padding), mode="reflect").squeeze(1)
|
||||
|
||||
spec = torch.stft(
|
||||
waveform,
|
||||
n_fft,
|
||||
hop_length=hop_length,
|
||||
win_length=win_length,
|
||||
window=hann_window,
|
||||
center=center,
|
||||
pad_mode="reflect",
|
||||
normalized=False,
|
||||
onesided=True,
|
||||
return_complex=True,
|
||||
)
|
||||
spec = torch.sqrt(torch.view_as_real(spec).pow(2).sum(-1) + 1e-9)
|
||||
|
||||
mel_spec = torch.matmul(mel_basis, spec)
|
||||
mel_spec = torch.log(torch.clamp(mel_spec, min=1e-5))
|
||||
|
||||
return mel_spec
|
||||
|
||||
|
||||
def get_vocos_mel_spectrogram(
|
||||
waveform,
|
||||
n_fft=1024,
|
||||
n_mel_channels=100,
|
||||
target_sample_rate=24000,
|
||||
hop_length=256,
|
||||
win_length=1024,
|
||||
):
|
||||
mel_stft = torchaudio.transforms.MelSpectrogram(
|
||||
sample_rate=target_sample_rate,
|
||||
n_fft=n_fft,
|
||||
win_length=win_length,
|
||||
hop_length=hop_length,
|
||||
n_mels=n_mel_channels,
|
||||
power=1,
|
||||
center=True,
|
||||
normalized=False,
|
||||
norm=None,
|
||||
).to(waveform.device)
|
||||
if len(waveform.shape) == 3:
|
||||
waveform = waveform.squeeze(1) # 'b 1 nw -> b nw'
|
||||
|
||||
assert len(waveform.shape) == 2
|
||||
|
||||
mel = mel_stft(waveform)
|
||||
mel = mel.clamp(min=1e-5).log()
|
||||
return mel
|
||||
|
||||
|
||||
class MelSpec(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
n_fft=1024,
|
||||
hop_length=256,
|
||||
win_length=1024,
|
||||
n_mel_channels=100,
|
||||
target_sample_rate=24_000,
|
||||
mel_spec_type="vocos",
|
||||
):
|
||||
super().__init__()
|
||||
assert mel_spec_type in ["vocos", "bigvgan"], print("We only support two extract mel backend: vocos or bigvgan")
|
||||
|
||||
self.n_fft = n_fft
|
||||
self.hop_length = hop_length
|
||||
self.win_length = win_length
|
||||
self.n_mel_channels = n_mel_channels
|
||||
self.target_sample_rate = target_sample_rate
|
||||
|
||||
if mel_spec_type == "vocos":
|
||||
self.extractor = get_vocos_mel_spectrogram
|
||||
elif mel_spec_type == "bigvgan":
|
||||
self.extractor = get_bigvgan_mel_spectrogram
|
||||
|
||||
self.register_buffer("dummy", torch.tensor(0), persistent=False)
|
||||
|
||||
def forward(self, wav):
|
||||
if self.dummy.device != wav.device:
|
||||
self.to(wav.device)
|
||||
|
||||
mel = self.extractor(
|
||||
waveform=wav,
|
||||
n_fft=self.n_fft,
|
||||
n_mel_channels=self.n_mel_channels,
|
||||
target_sample_rate=self.target_sample_rate,
|
||||
hop_length=self.hop_length,
|
||||
win_length=self.win_length,
|
||||
)
|
||||
|
||||
return mel
|
||||
|
||||
|
||||
# sinusoidal position embedding
|
||||
|
||||
|
||||
@@ -328,20 +202,18 @@ class FeedForward(nn.Module):
|
||||
return self.ff(x)
|
||||
|
||||
|
||||
# Attention with possible joint part
|
||||
# Attention
|
||||
# modified from diffusers/src/diffusers/models/attention_processor.py
|
||||
|
||||
|
||||
class Attention(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
processor: JointAttnProcessor | AttnProcessor,
|
||||
processor: AttnProcessor,
|
||||
dim: int,
|
||||
heads: int = 8,
|
||||
dim_head: int = 64,
|
||||
dropout: float = 0.0,
|
||||
context_dim: Optional[int] = None, # if not None -> joint attention
|
||||
context_pre_only=None,
|
||||
):
|
||||
super().__init__()
|
||||
|
||||
@@ -355,38 +227,21 @@ class Attention(nn.Module):
|
||||
self.inner_dim = dim_head * heads
|
||||
self.dropout = dropout
|
||||
|
||||
self.context_dim = context_dim
|
||||
self.context_pre_only = context_pre_only
|
||||
|
||||
self.to_q = nn.Linear(dim, self.inner_dim)
|
||||
self.to_k = nn.Linear(dim, self.inner_dim)
|
||||
self.to_v = nn.Linear(dim, self.inner_dim)
|
||||
|
||||
if self.context_dim is not None:
|
||||
self.to_k_c = nn.Linear(context_dim, self.inner_dim)
|
||||
self.to_v_c = nn.Linear(context_dim, self.inner_dim)
|
||||
if self.context_pre_only is not None:
|
||||
self.to_q_c = nn.Linear(context_dim, self.inner_dim)
|
||||
|
||||
self.to_out = nn.ModuleList([])
|
||||
self.to_out.append(nn.Linear(self.inner_dim, dim))
|
||||
self.to_out.append(nn.Dropout(dropout))
|
||||
|
||||
if self.context_pre_only is not None and not self.context_pre_only:
|
||||
self.to_out_c = nn.Linear(self.inner_dim, dim)
|
||||
|
||||
def forward(
|
||||
self,
|
||||
x: float["b n d"], # noised input x # noqa: F722
|
||||
c: float["b n d"] = None, # context c # noqa: F722
|
||||
mask: bool["b n"] | None = None, # noqa: F722
|
||||
rope=None, # rotary position embedding for x
|
||||
c_rope=None, # rotary position embedding for c
|
||||
) -> torch.Tensor:
|
||||
if c is not None:
|
||||
return self.processor(self, x, c=c, mask=mask, rope=rope, c_rope=c_rope)
|
||||
else:
|
||||
return self.processor(self, x, mask=mask, rope=rope)
|
||||
return self.processor(self, x, mask=mask, rope=rope)
|
||||
|
||||
|
||||
# Attention processor
|
||||
@@ -457,93 +312,6 @@ class AttnProcessor:
|
||||
return x
|
||||
|
||||
|
||||
# Joint Attention processor for MM-DiT
|
||||
# modified from diffusers/src/diffusers/models/attention_processor.py
|
||||
|
||||
|
||||
class JointAttnProcessor:
|
||||
def __init__(self):
|
||||
pass
|
||||
|
||||
def __call__(
|
||||
self,
|
||||
attn: Attention,
|
||||
x: float["b n d"], # noised input x # noqa: F722
|
||||
c: float["b nt d"] = None, # context c, here text # noqa: F722
|
||||
mask: bool["b n"] | None = None, # noqa: F722
|
||||
rope=None, # rotary position embedding for x
|
||||
c_rope=None, # rotary position embedding for c
|
||||
) -> torch.FloatTensor:
|
||||
residual = x
|
||||
|
||||
batch_size = c.shape[0]
|
||||
|
||||
# `sample` projections.
|
||||
query = attn.to_q(x)
|
||||
key = attn.to_k(x)
|
||||
value = attn.to_v(x)
|
||||
|
||||
# `context` projections.
|
||||
c_query = attn.to_q_c(c)
|
||||
c_key = attn.to_k_c(c)
|
||||
c_value = attn.to_v_c(c)
|
||||
|
||||
# apply rope for context and noised input independently
|
||||
if rope is not None:
|
||||
freqs, xpos_scale = rope
|
||||
q_xpos_scale, k_xpos_scale = (xpos_scale, xpos_scale**-1.0) if xpos_scale is not None else (1.0, 1.0)
|
||||
query = apply_rotary_pos_emb(query, freqs, q_xpos_scale)
|
||||
key = apply_rotary_pos_emb(key, freqs, k_xpos_scale)
|
||||
if c_rope is not None:
|
||||
freqs, xpos_scale = c_rope
|
||||
q_xpos_scale, k_xpos_scale = (xpos_scale, xpos_scale**-1.0) if xpos_scale is not None else (1.0, 1.0)
|
||||
c_query = apply_rotary_pos_emb(c_query, freqs, q_xpos_scale)
|
||||
c_key = apply_rotary_pos_emb(c_key, freqs, k_xpos_scale)
|
||||
|
||||
# attention
|
||||
query = torch.cat([query, c_query], dim=1)
|
||||
key = torch.cat([key, c_key], dim=1)
|
||||
value = torch.cat([value, c_value], dim=1)
|
||||
|
||||
inner_dim = key.shape[-1]
|
||||
head_dim = inner_dim // attn.heads
|
||||
query = query.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
|
||||
key = key.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
|
||||
value = value.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
|
||||
|
||||
# mask. e.g. inference got a batch with different target durations, mask out the padding
|
||||
if mask is not None:
|
||||
attn_mask = F.pad(mask, (0, c.shape[1]), value=True) # no mask for c (text)
|
||||
attn_mask = attn_mask.unsqueeze(1).unsqueeze(1) # 'b n -> b 1 1 n'
|
||||
attn_mask = attn_mask.expand(batch_size, attn.heads, query.shape[-2], key.shape[-2])
|
||||
else:
|
||||
attn_mask = None
|
||||
|
||||
x = F.scaled_dot_product_attention(query, key, value, attn_mask=attn_mask, dropout_p=0.0, is_causal=False)
|
||||
x = x.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim)
|
||||
x = x.to(query.dtype)
|
||||
|
||||
# Split the attention outputs.
|
||||
x, c = (
|
||||
x[:, : residual.shape[1]],
|
||||
x[:, residual.shape[1] :],
|
||||
)
|
||||
|
||||
# linear proj
|
||||
x = attn.to_out[0](x)
|
||||
# dropout
|
||||
x = attn.to_out[1](x)
|
||||
if not attn.context_pre_only:
|
||||
c = attn.to_out_c(c)
|
||||
|
||||
if mask is not None:
|
||||
mask = mask.unsqueeze(-1)
|
||||
x = x.masked_fill(~mask, 0.0)
|
||||
# c = c.masked_fill(~mask, 0.) # no mask for c (text)
|
||||
|
||||
return x, c
|
||||
|
||||
|
||||
# DiT Block
|
||||
|
||||
|
||||
@@ -580,76 +348,6 @@ class DiTBlock(nn.Module):
|
||||
return x
|
||||
|
||||
|
||||
# MMDiT Block https://arxiv.org/abs/2403.03206
|
||||
|
||||
|
||||
class MMDiTBlock(nn.Module):
|
||||
r"""
|
||||
modified from diffusers/src/diffusers/models/attention.py
|
||||
|
||||
notes.
|
||||
_c: context related. text, cond, etc. (left part in sd3 fig2.b)
|
||||
_x: noised input related. (right part)
|
||||
context_pre_only: last layer only do prenorm + modulation cuz no more ffn
|
||||
"""
|
||||
|
||||
def __init__(self, dim, heads, dim_head, ff_mult=4, dropout=0.1, context_pre_only=False):
|
||||
super().__init__()
|
||||
|
||||
self.context_pre_only = context_pre_only
|
||||
|
||||
self.attn_norm_c = AdaLayerNormZero_Final(dim) if context_pre_only else AdaLayerNormZero(dim)
|
||||
self.attn_norm_x = AdaLayerNormZero(dim)
|
||||
self.attn = Attention(
|
||||
processor=JointAttnProcessor(),
|
||||
dim=dim,
|
||||
heads=heads,
|
||||
dim_head=dim_head,
|
||||
dropout=dropout,
|
||||
context_dim=dim,
|
||||
context_pre_only=context_pre_only,
|
||||
)
|
||||
|
||||
if not context_pre_only:
|
||||
self.ff_norm_c = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-6)
|
||||
self.ff_c = FeedForward(dim=dim, mult=ff_mult, dropout=dropout, approximate="tanh")
|
||||
else:
|
||||
self.ff_norm_c = None
|
||||
self.ff_c = None
|
||||
self.ff_norm_x = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-6)
|
||||
self.ff_x = FeedForward(dim=dim, mult=ff_mult, dropout=dropout, approximate="tanh")
|
||||
|
||||
def forward(self, x, c, t, mask=None, rope=None, c_rope=None): # x: noised input, c: context, t: time embedding
|
||||
# pre-norm & modulation for attention input
|
||||
if self.context_pre_only:
|
||||
norm_c = self.attn_norm_c(c, t)
|
||||
else:
|
||||
norm_c, c_gate_msa, c_shift_mlp, c_scale_mlp, c_gate_mlp = self.attn_norm_c(c, emb=t)
|
||||
norm_x, x_gate_msa, x_shift_mlp, x_scale_mlp, x_gate_mlp = self.attn_norm_x(x, emb=t)
|
||||
|
||||
# attention
|
||||
x_attn_output, c_attn_output = self.attn(x=norm_x, c=norm_c, mask=mask, rope=rope, c_rope=c_rope)
|
||||
|
||||
# process attention output for context c
|
||||
if self.context_pre_only:
|
||||
c = None
|
||||
else: # if not last layer
|
||||
c = c + c_gate_msa.unsqueeze(1) * c_attn_output
|
||||
|
||||
norm_c = self.ff_norm_c(c) * (1 + c_scale_mlp[:, None]) + c_shift_mlp[:, None]
|
||||
c_ff_output = self.ff_c(norm_c)
|
||||
c = c + c_gate_mlp.unsqueeze(1) * c_ff_output
|
||||
|
||||
# process attention output for input x
|
||||
x = x + x_gate_msa.unsqueeze(1) * x_attn_output
|
||||
|
||||
norm_x = self.ff_norm_x(x) * (1 + x_scale_mlp[:, None]) + x_shift_mlp[:, None]
|
||||
x_ff_output = self.ff_x(norm_x)
|
||||
x = x + x_gate_mlp.unsqueeze(1) * x_ff_output
|
||||
|
||||
return c, x
|
||||
|
||||
|
||||
# time step conditioning embedding
|
||||
|
||||
|
||||
|
||||
+1432
-1337
File diff suppressed because it is too large
Load Diff
@@ -6,7 +6,20 @@
|
||||
全部按英文识别
|
||||
全部按日文识别
|
||||
"""
|
||||
import psutil
|
||||
import os
|
||||
|
||||
def set_high_priority():
|
||||
"""把当前 Python 进程设为 HIGH_PRIORITY_CLASS"""
|
||||
if os.name != "nt":
|
||||
return # 仅 Windows 有效
|
||||
p = psutil.Process(os.getpid())
|
||||
try:
|
||||
p.nice(psutil.HIGH_PRIORITY_CLASS)
|
||||
print("已将进程优先级设为 High")
|
||||
except psutil.AccessDenied:
|
||||
print("权限不足,无法修改优先级(请用管理员运行)")
|
||||
set_high_priority()
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
@@ -112,13 +125,14 @@ is_exist_s2gv4 = os.path.exists(path_sovits_v4)
|
||||
tts_config = TTS_Config("GPT_SoVITS/configs/tts_infer.yaml")
|
||||
tts_config.device = device
|
||||
tts_config.is_half = is_half
|
||||
tts_config.version = version
|
||||
# tts_config.version = version
|
||||
tts_config.update_version(version)
|
||||
if gpt_path is not None:
|
||||
if "!" in gpt_path:
|
||||
if "!" in gpt_path or "!" in gpt_path:
|
||||
gpt_path = name2gpt_path[gpt_path]
|
||||
tts_config.t2s_weights_path = gpt_path
|
||||
if sovits_path is not None:
|
||||
if "!" in sovits_path:
|
||||
if "!" in sovits_path or "!" in sovits_path:
|
||||
sovits_path = name2sovits_path[sovits_path]
|
||||
tts_config.vits_weights_path = sovits_path
|
||||
if cnhubert_base_path is not None:
|
||||
@@ -217,7 +231,7 @@ v3v4set = {"v3", "v4"}
|
||||
|
||||
|
||||
def change_sovits_weights(sovits_path, prompt_language=None, text_language=None):
|
||||
if "!" in sovits_path:
|
||||
if "!" in sovits_path or "!" in sovits_path:
|
||||
sovits_path = name2sovits_path[sovits_path]
|
||||
global version, model_version, dict_language, if_lora_v3
|
||||
version, model_version, if_lora_v3 = get_sovits_version_from_path_fast(sovits_path)
|
||||
@@ -283,6 +297,12 @@ def change_sovits_weights(sovits_path, prompt_language=None, text_language=None)
|
||||
f.write(json.dumps(data))
|
||||
|
||||
|
||||
def change_gpt_weights(gpt_path):
|
||||
if "!" in gpt_path or "!" in gpt_path:
|
||||
gpt_path = name2gpt_path[gpt_path]
|
||||
tts_pipeline.init_t2s_weights(gpt_path)
|
||||
|
||||
|
||||
with gr.Blocks(title="GPT-SoVITS WebUI", analytics_enabled=False, js=js, css=css) as app:
|
||||
gr.HTML(
|
||||
top_html.format(
|
||||
@@ -365,7 +385,7 @@ with gr.Blocks(title="GPT-SoVITS WebUI", analytics_enabled=False, js=js, css=css
|
||||
minimum=0.6, maximum=1.65, step=0.05, label="语速", value=1.0, interactive=True
|
||||
)
|
||||
with gr.Row():
|
||||
top_k = gr.Slider(minimum=1, maximum=100, step=1, label=i18n("top_k"), value=5, interactive=True)
|
||||
top_k = gr.Slider(minimum=1, maximum=100, step=1, label=i18n("top_k"), value=15, interactive=True)
|
||||
top_p = gr.Slider(minimum=0, maximum=1, step=0.05, label=i18n("top_p"), value=1, interactive=True)
|
||||
with gr.Row():
|
||||
temperature = gr.Slider(
|
||||
@@ -457,7 +477,7 @@ with gr.Blocks(title="GPT-SoVITS WebUI", analytics_enabled=False, js=js, css=css
|
||||
inference_button,
|
||||
],
|
||||
) #
|
||||
GPT_dropdown.change(tts_pipeline.init_t2s_weights, [GPT_dropdown], [])
|
||||
GPT_dropdown.change(change_gpt_weights, [GPT_dropdown], [])
|
||||
|
||||
with gr.Group():
|
||||
gr.Markdown(
|
||||
|
||||
@@ -37,6 +37,10 @@ from einops import rearrange, repeat
|
||||
import torch
|
||||
from torch import nn
|
||||
import torch.nn.functional as F
|
||||
import torch.distributed as dist
|
||||
|
||||
from module.distrib import broadcast_tensors, is_distributed
|
||||
from module.ddp_utils import SyncFunction
|
||||
from tqdm import tqdm
|
||||
|
||||
|
||||
@@ -69,27 +73,40 @@ def sample_vectors(samples, num: int):
|
||||
return samples[indices]
|
||||
|
||||
|
||||
def kmeans(samples, num_clusters: int, num_iters: int = 10):
|
||||
dim, dtype = samples.shape[-1], samples.dtype
|
||||
max_kmeans_samples = 500
|
||||
samples = samples[:max_kmeans_samples, :]
|
||||
def kmeans(samples, num_clusters: int, num_iters: int = 10, frames_to_use: int = 10_000, batch_size: int = 64):
|
||||
N, D = samples.shape
|
||||
dtype, device = samples.dtype, samples.device
|
||||
|
||||
if frames_to_use < N:
|
||||
indices = torch.randperm(N, device=device)[:frames_to_use]
|
||||
samples = samples[indices]
|
||||
|
||||
means = sample_vectors(samples, num_clusters)
|
||||
|
||||
print("kmeans start ... ")
|
||||
for _ in tqdm(range(num_iters)):
|
||||
diffs = rearrange(samples, "n d -> n () d") - rearrange(means, "c d -> () c d")
|
||||
dists = -(diffs**2).sum(dim=-1)
|
||||
# Store cluster assignments
|
||||
all_assignments = []
|
||||
|
||||
buckets = dists.max(dim=-1).indices
|
||||
for i in range(0, samples.shape[0], batch_size):
|
||||
batch = samples[i : i + batch_size] # [B, D]
|
||||
dists = torch.cdist(batch, means, p=2) # [B, C]
|
||||
assignments = dists.argmin(dim=1) # [B]
|
||||
all_assignments.append(assignments)
|
||||
|
||||
buckets = torch.cat(all_assignments, dim=0) # [N]
|
||||
bins = torch.bincount(buckets, minlength=num_clusters)
|
||||
zero_mask = bins == 0
|
||||
bins_min_clamped = bins.masked_fill(zero_mask, 1)
|
||||
|
||||
new_means = buckets.new_zeros(num_clusters, dim, dtype=dtype)
|
||||
new_means.scatter_add_(0, repeat(buckets, "n -> n d", d=dim), samples)
|
||||
new_means = new_means / bins_min_clamped[..., None]
|
||||
# Compute new means
|
||||
new_means = torch.zeros_like(means)
|
||||
for i in range(num_clusters):
|
||||
mask = buckets == i
|
||||
if mask.any():
|
||||
new_means[i] = samples[mask].mean(dim=0)
|
||||
|
||||
means = torch.where(zero_mask[..., None], means, new_means)
|
||||
means = torch.where(zero_mask[:, None], means, new_means)
|
||||
|
||||
return means, bins
|
||||
|
||||
@@ -141,13 +158,24 @@ class EuclideanCodebook(nn.Module):
|
||||
if self.inited:
|
||||
return
|
||||
|
||||
embed, cluster_size = kmeans(data, self.codebook_size, self.kmeans_iters)
|
||||
if dist.is_available() and dist.is_initialized():
|
||||
# [B * T * world_size, D]
|
||||
data = SyncFunction.apply(data)
|
||||
|
||||
if dist.get_rank() == 0:
|
||||
embed, cluster_size = kmeans(data, self.codebook_size, self.kmeans_iters)
|
||||
else:
|
||||
embed = torch.empty_like(self.embed)
|
||||
cluster_size = torch.empty_like(self.cluster_size)
|
||||
dist.broadcast(embed, src=0)
|
||||
dist.broadcast(cluster_size, src=0)
|
||||
|
||||
self.embed.data.copy_(embed)
|
||||
self.embed_avg.data.copy_(embed.clone())
|
||||
self.cluster_size.data.copy_(cluster_size)
|
||||
self.inited.data.copy_(torch.Tensor([True]))
|
||||
# Make sure all buffers across workers are in sync after initialization
|
||||
# broadcast_tensors(self.buffers())
|
||||
broadcast_tensors(self.buffers())
|
||||
|
||||
def replace_(self, samples, mask):
|
||||
modified_codebook = torch.where(mask[..., None], sample_vectors(samples, self.codebook_size), self.embed)
|
||||
@@ -161,9 +189,17 @@ class EuclideanCodebook(nn.Module):
|
||||
if not torch.any(expired_codes):
|
||||
return
|
||||
|
||||
batch_samples = rearrange(batch_samples, "... d -> (...) d")
|
||||
self.replace_(batch_samples, mask=expired_codes)
|
||||
# broadcast_tensors(self.buffers())
|
||||
if is_distributed():
|
||||
# [B * T * world_size, D]
|
||||
batch_samples = SyncFunction.apply(batch_samples)
|
||||
|
||||
if dist.get_rank() == 0:
|
||||
new_embeds = sample_vectors(batch_samples, expired_codes.sum())
|
||||
else:
|
||||
new_embeds = torch.zeros(expired_codes.sum(), self.embed.size(1), device=self.embed.device)
|
||||
dist.broadcast(new_embeds, src=0)
|
||||
self.embed.data[expired_codes] = new_embeds
|
||||
broadcast_tensors(self.buffers())
|
||||
|
||||
def preprocess(self, x):
|
||||
x = rearrange(x, "... d -> (...) d")
|
||||
@@ -208,17 +244,26 @@ class EuclideanCodebook(nn.Module):
|
||||
quantize = self.dequantize(embed_ind)
|
||||
|
||||
if self.training:
|
||||
### Update codebook by EMA
|
||||
embed_onehot_sum = embed_onehot.sum(0) # [cb-size,]
|
||||
embed_sum = x.t() @ embed_onehot # [D, cb-size]
|
||||
if is_distributed():
|
||||
dist.all_reduce(embed_onehot_sum)
|
||||
dist.all_reduce(embed_sum)
|
||||
# Update ema cluster count N_i^t, eq. (6) in vqvae paper
|
||||
self.cluster_size.data.mul_(self.decay).add_(embed_onehot_sum, alpha=1 - self.decay)
|
||||
# Update ema embed: eq. (7) in vqvae paper
|
||||
self.embed_avg.data.mul_(self.decay).add_(embed_sum.t(), alpha=1 - self.decay)
|
||||
# apply laplace smoothing
|
||||
n = self.cluster_size.sum()
|
||||
cluster_size = (self.cluster_size + self.epsilon) / (n + self.codebook_size * self.epsilon) * n
|
||||
# Update ema embed: eq. (8) in vqvae paper
|
||||
embed_normalized = self.embed_avg / cluster_size.unsqueeze(1)
|
||||
self.embed.data.copy_(embed_normalized)
|
||||
|
||||
# We do the expiry of code at that point as buffers are in sync
|
||||
# and all the workers will take the same decision.
|
||||
self.expire_codes_(x)
|
||||
ema_inplace(self.cluster_size, embed_onehot.sum(0), self.decay)
|
||||
embed_sum = x.t() @ embed_onehot
|
||||
ema_inplace(self.embed_avg, embed_sum.t(), self.decay)
|
||||
cluster_size = (
|
||||
laplace_smoothing(self.cluster_size, self.codebook_size, self.epsilon) * self.cluster_size.sum()
|
||||
)
|
||||
embed_normalized = self.embed_avg / cluster_size.unsqueeze(1)
|
||||
self.embed.data.copy_(embed_normalized)
|
||||
|
||||
return quantize, embed_ind
|
||||
|
||||
|
||||
@@ -21,7 +21,7 @@ class TextAudioSpeakerLoader(torch.utils.data.Dataset):
|
||||
3) computes spectrograms from audio files.
|
||||
"""
|
||||
|
||||
def __init__(self, hparams, version=None,val=False):
|
||||
def __init__(self, hparams, version=None, val=False):
|
||||
exp_dir = hparams.exp_dir
|
||||
self.path2 = "%s/2-name2text.txt" % exp_dir
|
||||
self.path4 = "%s/4-cnhubert" % exp_dir
|
||||
@@ -29,7 +29,7 @@ class TextAudioSpeakerLoader(torch.utils.data.Dataset):
|
||||
assert os.path.exists(self.path2)
|
||||
assert os.path.exists(self.path4)
|
||||
assert os.path.exists(self.path5)
|
||||
self.is_v2Pro=version in {"v2Pro","v2ProPlus"}
|
||||
self.is_v2Pro = version in {"v2Pro", "v2ProPlus"}
|
||||
if self.is_v2Pro:
|
||||
self.path7 = "%s/7-sv_cn" % exp_dir
|
||||
assert os.path.exists(self.path7)
|
||||
@@ -118,7 +118,7 @@ class TextAudioSpeakerLoader(torch.utils.data.Dataset):
|
||||
ssl = F.pad(ssl.float(), (0, 1), mode="replicate").to(typee)
|
||||
ssl.requires_grad = False
|
||||
if self.is_v2Pro:
|
||||
sv_emb=torch.load("%s/%s.pt" % (self.path7, audiopath), map_location="cpu")
|
||||
sv_emb = torch.load("%s/%s.pt" % (self.path7, audiopath), map_location="cpu")
|
||||
except:
|
||||
traceback.print_exc()
|
||||
spec = torch.zeros(1025, 100)
|
||||
@@ -126,10 +126,10 @@ class TextAudioSpeakerLoader(torch.utils.data.Dataset):
|
||||
ssl = torch.zeros(1, 768, 100)
|
||||
text = text[-1:]
|
||||
if self.is_v2Pro:
|
||||
sv_emb=torch.zeros(1,20480)
|
||||
sv_emb = torch.zeros(1, 20480)
|
||||
print("load audio or ssl error!!!!!!", audiopath)
|
||||
if self.is_v2Pro:
|
||||
return (ssl, spec, wav, text,sv_emb)
|
||||
return (ssl, spec, wav, text, sv_emb)
|
||||
else:
|
||||
return (ssl, spec, wav, text)
|
||||
|
||||
@@ -192,9 +192,9 @@ class TextAudioSpeakerLoader(torch.utils.data.Dataset):
|
||||
class TextAudioSpeakerCollate:
|
||||
"""Zero-pads model inputs and targets"""
|
||||
|
||||
def __init__(self, return_ids=False,version=None):
|
||||
def __init__(self, return_ids=False, version=None):
|
||||
self.return_ids = return_ids
|
||||
self.is_v2Pro=version in {"v2Pro","v2ProPlus"}
|
||||
self.is_v2Pro = version in {"v2Pro", "v2ProPlus"}
|
||||
|
||||
def __call__(self, batch):
|
||||
"""Collate's training batch from normalized text, audio and speaker identities
|
||||
@@ -228,7 +228,7 @@ class TextAudioSpeakerCollate:
|
||||
text_padded.zero_()
|
||||
|
||||
if self.is_v2Pro:
|
||||
sv_embs=torch.FloatTensor(len(batch),20480)
|
||||
sv_embs = torch.FloatTensor(len(batch), 20480)
|
||||
|
||||
for i in range(len(ids_sorted_decreasing)):
|
||||
row = batch[ids_sorted_decreasing[i]]
|
||||
@@ -250,11 +250,30 @@ class TextAudioSpeakerCollate:
|
||||
text_lengths[i] = text.size(0)
|
||||
|
||||
if self.is_v2Pro:
|
||||
sv_embs[i]=row[4]
|
||||
sv_embs[i] = row[4]
|
||||
if self.is_v2Pro:
|
||||
return ssl_padded, ssl_lengths, spec_padded, spec_lengths, wav_padded, wav_lengths, text_padded, text_lengths,sv_embs
|
||||
return (
|
||||
ssl_padded,
|
||||
ssl_lengths,
|
||||
spec_padded,
|
||||
spec_lengths,
|
||||
wav_padded,
|
||||
wav_lengths,
|
||||
text_padded,
|
||||
text_lengths,
|
||||
sv_embs,
|
||||
)
|
||||
else:
|
||||
return ssl_padded, ssl_lengths, spec_padded, spec_lengths, wav_padded, wav_lengths, text_padded, text_lengths
|
||||
return (
|
||||
ssl_padded,
|
||||
ssl_lengths,
|
||||
spec_padded,
|
||||
spec_lengths,
|
||||
wav_padded,
|
||||
wav_lengths,
|
||||
text_padded,
|
||||
text_lengths,
|
||||
)
|
||||
|
||||
|
||||
class TextAudioSpeakerLoaderV3(torch.utils.data.Dataset):
|
||||
|
||||
@@ -0,0 +1,181 @@
|
||||
import torch
|
||||
from torch.nn.parallel import DistributedDataParallel
|
||||
from torch.nn.parallel.distributed import _find_tensors
|
||||
from packaging import version
|
||||
|
||||
|
||||
# from https://github.com/Lightning-AI/lightning-bolts/blob/5d61197cd2f491f69e238137a5edabe80ae14ad9/pl_bolts/models/self_supervised/simclr/simclr_module.py#L20
|
||||
class SyncFunction(torch.autograd.Function):
|
||||
@staticmethod
|
||||
# @torch.no_grad()
|
||||
def forward(ctx, tensor):
|
||||
world_size = torch.distributed.get_world_size()
|
||||
|
||||
# Collect batch sizes from all processes
|
||||
local_bs = torch.tensor([tensor.shape[0]], device=tensor.device)
|
||||
batch_sizes = [torch.zeros_like(local_bs) for _ in range(world_size)]
|
||||
torch.distributed.all_gather(batch_sizes, local_bs)
|
||||
|
||||
# Convert to integer list and find the minimum
|
||||
batch_sizes_int = [bs.item() for bs in batch_sizes]
|
||||
min_bs = min(batch_sizes_int)
|
||||
|
||||
# Crop the tensor to the minimum batch size if needed
|
||||
cropped_tensor = tensor[:min_bs] if tensor.shape[0] > min_bs else tensor
|
||||
|
||||
# Prepare for gathering
|
||||
out_shape = (min_bs * world_size,) + tensor.shape[1:]
|
||||
gathered_tensor = torch.zeros(out_shape, dtype=tensor.dtype, device=tensor.device)
|
||||
|
||||
# Build tensor list for all_gather
|
||||
tensor_list = list(torch.chunk(gathered_tensor, world_size))
|
||||
|
||||
# Perform all_gather using the cropped tensors
|
||||
torch.distributed.all_gather(tensor_list, cropped_tensor)
|
||||
|
||||
# Save for backward pass
|
||||
ctx.min_bs = min_bs
|
||||
ctx.world_size = world_size
|
||||
ctx.orig_shape = tensor.shape
|
||||
|
||||
return gathered_tensor
|
||||
|
||||
@staticmethod
|
||||
def backward(ctx, grad_output):
|
||||
assert False
|
||||
grad_input = grad_output.clone()
|
||||
torch.distributed.all_reduce(grad_input, op=torch.distributed.ReduceOp.SUM, async_op=False)
|
||||
|
||||
idx_from = torch.distributed.get_rank() * ctx.batch_size
|
||||
idx_to = (torch.distributed.get_rank() + 1) * ctx.batch_size
|
||||
return grad_input[idx_from:idx_to]
|
||||
|
||||
class DDP(DistributedDataParallel):
|
||||
"""
|
||||
Override the forward call in lightning so it goes to training and validation step respectively
|
||||
"""
|
||||
|
||||
def forward(self, *inputs, **kwargs): # pragma: no cover
|
||||
if version.parse(torch.__version__[:6]) < version.parse("1.11"):
|
||||
self._sync_params()
|
||||
inputs, kwargs = self.scatter(inputs, kwargs, self.device_ids)
|
||||
assert len(self.device_ids) == 1
|
||||
if self.module.training:
|
||||
output = self.module.training_step(*inputs[0], **kwargs[0])
|
||||
elif self.module.testing:
|
||||
output = self.module.test_step(*inputs[0], **kwargs[0])
|
||||
else:
|
||||
output = self.module.validation_step(*inputs[0], **kwargs[0])
|
||||
if torch.is_grad_enabled():
|
||||
# We'll return the output object verbatim since it is a freeform
|
||||
# object. We need to find any tensors in this object, though,
|
||||
# because we need to figure out which parameters were used during
|
||||
# this forward pass, to ensure we short circuit reduction for any
|
||||
# unused parameters. Only if `find_unused_parameters` is set.
|
||||
if self.find_unused_parameters:
|
||||
self.reducer.prepare_for_backward(list(_find_tensors(output)))
|
||||
else:
|
||||
self.reducer.prepare_for_backward([])
|
||||
else:
|
||||
from torch.nn.parallel.distributed import (
|
||||
Join,
|
||||
_DDPSink,
|
||||
_tree_flatten_with_rref,
|
||||
_tree_unflatten_with_rref,
|
||||
)
|
||||
|
||||
with torch.autograd.profiler.record_function("DistributedDataParallel.forward"):
|
||||
if torch.is_grad_enabled() and self.require_backward_grad_sync:
|
||||
self.logger.set_runtime_stats_and_log()
|
||||
self.num_iterations += 1
|
||||
self.reducer.prepare_for_forward()
|
||||
|
||||
# Notify the join context that this process has not joined, if
|
||||
# needed
|
||||
work = Join.notify_join_context(self)
|
||||
if work:
|
||||
self.reducer._set_forward_pass_work_handle(work, self._divide_by_initial_world_size)
|
||||
|
||||
# Calling _rebuild_buckets before forward compuation,
|
||||
# It may allocate new buckets before deallocating old buckets
|
||||
# inside _rebuild_buckets. To save peak memory usage,
|
||||
# call _rebuild_buckets before the peak memory usage increases
|
||||
# during forward computation.
|
||||
# This should be called only once during whole training period.
|
||||
if torch.is_grad_enabled() and self.reducer._rebuild_buckets():
|
||||
print("Reducer buckets have been rebuilt in this iteration.")
|
||||
self._has_rebuilt_buckets = True
|
||||
|
||||
# sync params according to location (before/after forward) user
|
||||
# specified as part of hook, if hook was specified.
|
||||
buffer_hook_registered = hasattr(self, "buffer_hook")
|
||||
if self._check_sync_bufs_pre_fwd():
|
||||
self._sync_buffers()
|
||||
|
||||
if self._join_config.enable:
|
||||
# Notify joined ranks whether they should sync in backwards pass or not.
|
||||
self._check_global_requires_backward_grad_sync(is_joined_rank=False)
|
||||
|
||||
inputs, kwargs = self.scatter(inputs, kwargs, self.device_ids)
|
||||
if self.module.training:
|
||||
output = self.module.training_step(*inputs[0], **kwargs[0])
|
||||
elif self.module.testing:
|
||||
output = self.module.test_step(*inputs[0], **kwargs[0])
|
||||
else:
|
||||
output = self.module.validation_step(*inputs[0], **kwargs[0])
|
||||
|
||||
# sync params according to location (before/after forward) user
|
||||
# specified as part of hook, if hook was specified.
|
||||
if self._check_sync_bufs_post_fwd():
|
||||
self._sync_buffers()
|
||||
|
||||
if torch.is_grad_enabled() and self.require_backward_grad_sync:
|
||||
self.require_forward_param_sync = True
|
||||
# We'll return the output object verbatim since it is a freeform
|
||||
# object. We need to find any tensors in this object, though,
|
||||
# because we need to figure out which parameters were used during
|
||||
# this forward pass, to ensure we short circuit reduction for any
|
||||
# unused parameters. Only if `find_unused_parameters` is set.
|
||||
if self.find_unused_parameters and not self.static_graph:
|
||||
# Do not need to populate this for static graph.
|
||||
self.reducer.prepare_for_backward(list(_find_tensors(output)))
|
||||
else:
|
||||
self.reducer.prepare_for_backward([])
|
||||
else:
|
||||
self.require_forward_param_sync = False
|
||||
|
||||
# TODO: DDPSink is currently enabled for unused parameter detection and
|
||||
# static graph training for first iteration.
|
||||
if (self.find_unused_parameters and not self.static_graph) or (
|
||||
self.static_graph and self.num_iterations == 1
|
||||
):
|
||||
state_dict = {
|
||||
"static_graph": self.static_graph,
|
||||
"num_iterations": self.num_iterations,
|
||||
}
|
||||
|
||||
output_tensor_list, treespec, output_is_rref = _tree_flatten_with_rref(output)
|
||||
output_placeholders = [None for _ in range(len(output_tensor_list))]
|
||||
# Do not touch tensors that have no grad_fn, which can cause issues
|
||||
# such as https://github.com/pytorch/pytorch/issues/60733
|
||||
for i, output in enumerate(output_tensor_list):
|
||||
if torch.is_tensor(output) and output.grad_fn is None:
|
||||
output_placeholders[i] = output
|
||||
|
||||
# When find_unused_parameters=True, makes tensors which require grad
|
||||
# run through the DDPSink backward pass. When not all outputs are
|
||||
# used in loss, this makes those corresponding tensors receive
|
||||
# undefined gradient which the reducer then handles to ensure
|
||||
# param.grad field is not touched and we don't error out.
|
||||
passthrough_tensor_list = _DDPSink.apply(
|
||||
self.reducer,
|
||||
state_dict,
|
||||
*output_tensor_list,
|
||||
)
|
||||
for i in range(len(output_placeholders)):
|
||||
if output_placeholders[i] is None:
|
||||
output_placeholders[i] = passthrough_tensor_list[i]
|
||||
|
||||
# Reconstruct output data structure.
|
||||
output = _tree_unflatten_with_rref(output_placeholders, treespec, output_is_rref)
|
||||
return output
|
||||
@@ -0,0 +1,123 @@
|
||||
# Copyright (c) Meta Platforms, Inc. and affiliates.
|
||||
# All rights reserved.
|
||||
#
|
||||
# This source code is licensed under the license found in the
|
||||
# LICENSE file in the root directory of this source tree.
|
||||
|
||||
"""Torch distributed utilities."""
|
||||
|
||||
import typing as tp
|
||||
|
||||
import torch
|
||||
|
||||
|
||||
def rank():
|
||||
if torch.distributed.is_initialized():
|
||||
return torch.distributed.get_rank()
|
||||
else:
|
||||
return 0
|
||||
|
||||
|
||||
def world_size():
|
||||
if torch.distributed.is_initialized():
|
||||
return torch.distributed.get_world_size()
|
||||
else:
|
||||
return 1
|
||||
|
||||
|
||||
def is_distributed():
|
||||
return world_size() > 1
|
||||
|
||||
|
||||
def all_reduce(tensor: torch.Tensor, op=torch.distributed.ReduceOp.SUM):
|
||||
if is_distributed():
|
||||
return torch.distributed.all_reduce(tensor, op)
|
||||
|
||||
|
||||
def _is_complex_or_float(tensor):
|
||||
return torch.is_floating_point(tensor) or torch.is_complex(tensor)
|
||||
|
||||
|
||||
def _check_number_of_params(params: tp.List[torch.Tensor]):
|
||||
# utility function to check that the number of params in all workers is the same,
|
||||
# and thus avoid a deadlock with distributed all reduce.
|
||||
if not is_distributed() or not params:
|
||||
return
|
||||
# print('params[0].device ', params[0].device)
|
||||
tensor = torch.tensor([len(params)], device=params[0].device, dtype=torch.long)
|
||||
all_reduce(tensor)
|
||||
if tensor.item() != len(params) * world_size():
|
||||
# If not all the workers have the same number, for at least one of them,
|
||||
# this inequality will be verified.
|
||||
raise RuntimeError(
|
||||
f"Mismatch in number of params: ours is {len(params)}, at least one worker has a different one."
|
||||
)
|
||||
|
||||
|
||||
def broadcast_tensors(tensors: tp.Iterable[torch.Tensor], src: int = 0):
|
||||
"""Broadcast the tensors from the given parameters to all workers.
|
||||
This can be used to ensure that all workers have the same model to start with.
|
||||
"""
|
||||
if not is_distributed():
|
||||
return
|
||||
tensors = [tensor for tensor in tensors if _is_complex_or_float(tensor)]
|
||||
_check_number_of_params(tensors)
|
||||
handles = []
|
||||
for tensor in tensors:
|
||||
handle = torch.distributed.broadcast(tensor.data, src=src, async_op=True)
|
||||
handles.append(handle)
|
||||
for handle in handles:
|
||||
handle.wait()
|
||||
|
||||
|
||||
def sync_buffer(buffers, average=True):
|
||||
"""
|
||||
Sync grad for buffers. If average is False, broadcast instead of averaging.
|
||||
"""
|
||||
if not is_distributed():
|
||||
return
|
||||
handles = []
|
||||
for buffer in buffers:
|
||||
if torch.is_floating_point(buffer.data):
|
||||
if average:
|
||||
handle = torch.distributed.all_reduce(buffer.data, op=torch.distributed.ReduceOp.SUM, async_op=True)
|
||||
else:
|
||||
handle = torch.distributed.broadcast(buffer.data, src=0, async_op=True)
|
||||
handles.append((buffer, handle))
|
||||
for buffer, handle in handles:
|
||||
handle.wait()
|
||||
if average:
|
||||
buffer.data /= world_size()
|
||||
|
||||
|
||||
def sync_grad(params):
|
||||
"""
|
||||
Simpler alternative to DistributedDataParallel, that doesn't rely
|
||||
on any black magic. For simple models it can also be as fast.
|
||||
Just call this on your model parameters after the call to backward!
|
||||
"""
|
||||
if not is_distributed():
|
||||
return
|
||||
handles = []
|
||||
for p in params:
|
||||
if p.grad is not None:
|
||||
handle = torch.distributed.all_reduce(p.grad.data, op=torch.distributed.ReduceOp.SUM, async_op=True)
|
||||
handles.append((p, handle))
|
||||
for p, handle in handles:
|
||||
handle.wait()
|
||||
p.grad.data /= world_size()
|
||||
|
||||
|
||||
def average_metrics(metrics: tp.Dict[str, float], count=1.0):
|
||||
"""Average a dictionary of metrics across all workers, using the optional
|
||||
`count` as unormalized weight.
|
||||
"""
|
||||
if not is_distributed():
|
||||
return metrics
|
||||
keys, values = zip(*metrics.items())
|
||||
device = "cuda" if torch.cuda.is_available() else "cpu"
|
||||
tensor = torch.tensor(list(values) + [1], device=device, dtype=torch.float32)
|
||||
tensor *= count
|
||||
all_reduce(tensor)
|
||||
averaged = (tensor[:-1] / tensor[-1]).cpu().tolist()
|
||||
return dict(zip(keys, averaged))
|
||||
+143
-31
@@ -151,6 +151,8 @@ class DurationPredictor(nn.Module):
|
||||
return x * x_mask
|
||||
|
||||
|
||||
WINDOW = {}
|
||||
|
||||
class TextEncoder(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
@@ -209,7 +211,7 @@ class TextEncoder(nn.Module):
|
||||
|
||||
self.proj = nn.Conv1d(hidden_channels, out_channels * 2, 1)
|
||||
|
||||
def forward(self, y, y_lengths, text, text_lengths, ge, speed=1, test=None):
|
||||
def forward(self, y, y_lengths, text, text_lengths, ge, speed=1, test=None, result_length:int=None, overlap_frames:torch.Tensor=None, padding_length:int=None):
|
||||
y_mask = torch.unsqueeze(commons.sequence_mask(y_lengths, y.size(2)), 1).to(y.dtype)
|
||||
|
||||
y = self.ssl_proj(y * y_mask) * y_mask
|
||||
@@ -222,13 +224,44 @@ class TextEncoder(nn.Module):
|
||||
text = self.text_embedding(text).transpose(1, 2)
|
||||
text = self.encoder_text(text * text_mask, text_mask)
|
||||
y = self.mrte(y, y_mask, text, text_mask, ge)
|
||||
|
||||
if padding_length is not None and padding_length!=0:
|
||||
y = y[:, :, :-padding_length]
|
||||
y_mask = y_mask[:, :, :-padding_length]
|
||||
|
||||
|
||||
y = self.encoder2(y * y_mask, y_mask)
|
||||
|
||||
if result_length is not None:
|
||||
y = y[:, :, -result_length:]
|
||||
y_mask = y_mask[:, :, -result_length:]
|
||||
|
||||
if overlap_frames is not None:
|
||||
overlap_len = overlap_frames.shape[-1]
|
||||
window = WINDOW.get(overlap_len, None)
|
||||
if window is None:
|
||||
# WINDOW[overlap_len] = torch.hann_window(overlap_len*2, device=y.device, dtype=y.dtype)
|
||||
WINDOW[overlap_len] = torch.sin(torch.arange(overlap_len*2, device=y.device) * torch.pi / (overlap_len*2))
|
||||
window = WINDOW[overlap_len]
|
||||
|
||||
|
||||
window = window.to(y.device)
|
||||
y[:,:,:overlap_len] = (
|
||||
window[:overlap_len].view(1, 1, -1) * y[:,:,:overlap_len]
|
||||
+ window[overlap_len:].view(1, 1, -1) * overlap_frames
|
||||
)
|
||||
|
||||
y_ = y
|
||||
y_mask_ = y_mask
|
||||
|
||||
|
||||
|
||||
if speed != 1:
|
||||
y = F.interpolate(y, size=int(y.shape[-1] / speed) + 1, mode="linear")
|
||||
y_mask = F.interpolate(y_mask, size=y.shape[-1], mode="nearest")
|
||||
stats = self.proj(y) * y_mask
|
||||
m, logs = torch.split(stats, self.out_channels, dim=1)
|
||||
return y, m, logs, y_mask
|
||||
return y, m, logs, y_mask, y_, y_mask_
|
||||
|
||||
def extract_latent(self, x):
|
||||
x = self.ssl_proj(x)
|
||||
@@ -586,12 +619,17 @@ class DiscriminatorS(torch.nn.Module):
|
||||
|
||||
return x, fmap
|
||||
|
||||
v2pro_set={"v2Pro","v2ProPlus"}
|
||||
|
||||
v2pro_set = {"v2Pro", "v2ProPlus"}
|
||||
|
||||
|
||||
class MultiPeriodDiscriminator(torch.nn.Module):
|
||||
def __init__(self, use_spectral_norm=False,version=None):
|
||||
def __init__(self, use_spectral_norm=False, version=None):
|
||||
super(MultiPeriodDiscriminator, self).__init__()
|
||||
if version in v2pro_set:periods = [2, 3, 5, 7, 11,17,23]
|
||||
else:periods = [2, 3, 5, 7, 11]
|
||||
if version in v2pro_set:
|
||||
periods = [2, 3, 5, 7, 11, 17, 23]
|
||||
else:
|
||||
periods = [2, 3, 5, 7, 11]
|
||||
|
||||
discs = [DiscriminatorS(use_spectral_norm=use_spectral_norm)]
|
||||
discs = discs + [DiscriminatorP(i, use_spectral_norm=use_spectral_norm) for i in periods]
|
||||
@@ -787,6 +825,7 @@ class CodePredictor(nn.Module):
|
||||
|
||||
return pred_codes.transpose(0, 1)
|
||||
|
||||
|
||||
class SynthesizerTrn(nn.Module):
|
||||
"""
|
||||
Synthesizer for Training
|
||||
@@ -886,13 +925,13 @@ class SynthesizerTrn(nn.Module):
|
||||
self.quantizer = ResidualVectorQuantizer(dimension=ssl_dim, n_q=1, bins=1024)
|
||||
self.freeze_quantizer = freeze_quantizer
|
||||
|
||||
self.is_v2pro=self.version in v2pro_set
|
||||
self.is_v2pro = self.version in v2pro_set
|
||||
if self.is_v2pro:
|
||||
self.sv_emb = nn.Linear(20480, gin_channels)
|
||||
self.ge_to512 = nn.Linear(gin_channels, 512)
|
||||
self.prelu = nn.PReLU(num_parameters=gin_channels)
|
||||
|
||||
def forward(self, ssl, y, y_lengths, text, text_lengths,sv_emb=None):
|
||||
def forward(self, ssl, y, y_lengths, text, text_lengths, sv_emb=None):
|
||||
y_mask = torch.unsqueeze(commons.sequence_mask(y_lengths, y.size(2)), 1).to(y.dtype)
|
||||
if self.version == "v1":
|
||||
ge = self.ref_enc(y * y_mask, y_mask)
|
||||
@@ -915,7 +954,7 @@ class SynthesizerTrn(nn.Module):
|
||||
if self.semantic_frame_rate == "25hz":
|
||||
quantized = F.interpolate(quantized, size=int(quantized.shape[-1] * 2), mode="nearest")
|
||||
|
||||
x, m_p, logs_p, y_mask = self.enc_p(quantized, y_lengths, text, text_lengths, ge512 if self.is_v2pro else ge)
|
||||
x, m_p, logs_p, y_mask, _, _ = self.enc_p(quantized, y_lengths, text, text_lengths, ge512 if self.is_v2pro else ge)
|
||||
z, m_q, logs_q, y_mask = self.enc_q(y, y_lengths, g=ge)
|
||||
z_p = self.flow(z, y_mask, g=ge)
|
||||
|
||||
@@ -943,7 +982,7 @@ class SynthesizerTrn(nn.Module):
|
||||
if self.semantic_frame_rate == "25hz":
|
||||
quantized = F.interpolate(quantized, size=int(quantized.shape[-1] * 2), mode="nearest")
|
||||
|
||||
x, m_p, logs_p, y_mask = self.enc_p(quantized, y_lengths, text, text_lengths, ge, test=test)
|
||||
x, m_p, logs_p, y_mask, _, _ = self.enc_p(quantized, y_lengths, text, text_lengths, ge, test=test)
|
||||
z_p = m_p + torch.randn_like(m_p) * torch.exp(logs_p) * noise_scale
|
||||
|
||||
z = self.flow(z_p, y_mask, g=ge, reverse=True)
|
||||
@@ -951,8 +990,9 @@ class SynthesizerTrn(nn.Module):
|
||||
o = self.dec((z * y_mask)[:, :, :], g=ge)
|
||||
return o, y_mask, (z, z_p, m_p, logs_p)
|
||||
|
||||
|
||||
@torch.no_grad()
|
||||
def decode(self, codes, text, refer,noise_scale=0.5, speed=1, sv_emb=None):
|
||||
def decode(self, codes, text, refer, noise_scale=0.5, speed=1, sv_emb=None):
|
||||
def get_ge(refer, sv_emb):
|
||||
ge = None
|
||||
if refer is not None:
|
||||
@@ -970,8 +1010,8 @@ class SynthesizerTrn(nn.Module):
|
||||
|
||||
if type(refer) == list:
|
||||
ges = []
|
||||
for idx,_refer in enumerate(refer):
|
||||
ge = get_ge(_refer, sv_emb[idx]if self.is_v2pro else None)
|
||||
for idx, _refer in enumerate(refer):
|
||||
ge = get_ge(_refer, sv_emb[idx] if self.is_v2pro else None)
|
||||
ges.append(ge)
|
||||
ge = torch.stack(ges, 0).mean(0)
|
||||
else:
|
||||
@@ -983,7 +1023,14 @@ class SynthesizerTrn(nn.Module):
|
||||
quantized = self.quantizer.decode(codes)
|
||||
if self.semantic_frame_rate == "25hz":
|
||||
quantized = F.interpolate(quantized, size=int(quantized.shape[-1] * 2), mode="nearest")
|
||||
x, m_p, logs_p, y_mask = self.enc_p(quantized, y_lengths, text, text_lengths, self.ge_to512(ge.transpose(2,1)).transpose(2,1)if self.is_v2pro else ge, speed)
|
||||
x, m_p, logs_p, y_mask, _, _ = self.enc_p(
|
||||
quantized,
|
||||
y_lengths,
|
||||
text,
|
||||
text_lengths,
|
||||
self.ge_to512(ge.transpose(2, 1)).transpose(2, 1) if self.is_v2pro else ge,
|
||||
speed,
|
||||
)
|
||||
z_p = m_p + torch.randn_like(m_p) * torch.exp(logs_p) * noise_scale
|
||||
|
||||
z = self.flow(z_p, y_mask, g=ge, reverse=True)
|
||||
@@ -991,11 +1038,65 @@ class SynthesizerTrn(nn.Module):
|
||||
o = self.dec((z * y_mask)[:, :, :], g=ge)
|
||||
return o
|
||||
|
||||
|
||||
@torch.no_grad()
|
||||
def decode_streaming(self, codes, text, refer, noise_scale=0.5, speed=1, sv_emb=None, result_length:int=None, overlap_frames:torch.Tensor=None, padding_length:int=None):
|
||||
def get_ge(refer, sv_emb):
|
||||
ge = None
|
||||
if refer is not None:
|
||||
refer_lengths = torch.LongTensor([refer.size(2)]).to(refer.device)
|
||||
refer_mask = torch.unsqueeze(commons.sequence_mask(refer_lengths, refer.size(2)), 1).to(refer.dtype)
|
||||
if self.version == "v1":
|
||||
ge = self.ref_enc(refer * refer_mask, refer_mask)
|
||||
else:
|
||||
ge = self.ref_enc(refer[:, :704] * refer_mask, refer_mask)
|
||||
if self.is_v2pro:
|
||||
sv_emb = self.sv_emb(sv_emb) # B*20480->B*512
|
||||
ge += sv_emb.unsqueeze(-1)
|
||||
ge = self.prelu(ge)
|
||||
return ge
|
||||
|
||||
if type(refer) == list:
|
||||
ges = []
|
||||
for idx, _refer in enumerate(refer):
|
||||
ge = get_ge(_refer, sv_emb[idx] if self.is_v2pro else None)
|
||||
ges.append(ge)
|
||||
ge = torch.stack(ges, 0).mean(0)
|
||||
else:
|
||||
ge = get_ge(refer, sv_emb)
|
||||
|
||||
y_lengths = torch.LongTensor([codes.size(2) * 2]).to(codes.device)
|
||||
text_lengths = torch.LongTensor([text.size(-1)]).to(text.device)
|
||||
|
||||
quantized = self.quantizer.decode(codes)
|
||||
if self.semantic_frame_rate == "25hz":
|
||||
quantized = F.interpolate(quantized, size=int(quantized.shape[-1] * 2), mode="nearest")
|
||||
result_length = (2*result_length) if result_length is not None else None
|
||||
padding_length = (2*padding_length) if padding_length is not None else None
|
||||
x, m_p, logs_p, y_mask, y_, y_mask_ = self.enc_p(
|
||||
quantized,
|
||||
y_lengths,
|
||||
text,
|
||||
text_lengths,
|
||||
self.ge_to512(ge.transpose(2, 1)).transpose(2, 1) if self.is_v2pro else ge,
|
||||
speed,
|
||||
result_length=result_length,
|
||||
overlap_frames=overlap_frames,
|
||||
padding_length=padding_length
|
||||
)
|
||||
z_p = m_p + torch.randn_like(m_p) * torch.exp(logs_p) * noise_scale
|
||||
|
||||
z = self.flow(z_p, y_mask, g=ge, reverse=True)
|
||||
|
||||
o = self.dec((z * y_mask)[:, :, :], g=ge)
|
||||
return o, y_, y_mask_
|
||||
|
||||
def extract_latent(self, x):
|
||||
ssl = self.ssl_proj(x)
|
||||
quantized, codes, commit_loss, quantized_list = self.quantizer(ssl)
|
||||
return codes.transpose(0, 1)
|
||||
|
||||
|
||||
class CFM(torch.nn.Module):
|
||||
def __init__(self, in_channels, dit):
|
||||
super().__init__()
|
||||
@@ -1029,7 +1130,18 @@ class CFM(torch.nn.Module):
|
||||
t_tensor = torch.ones(x.shape[0], device=x.device, dtype=mu.dtype) * t
|
||||
# v_pred = model(x, t_tensor, d_tensor, **extra_args)
|
||||
v_pred, text_emb, dt = self.estimator(
|
||||
x, prompt_x, x_lens, t_tensor, d_tensor, mu, use_grad_ckpt=False, drop_audio_cond=False, drop_text=False, infer=True, text_cache=text_cache, dt_cache=dt_cache
|
||||
x,
|
||||
prompt_x,
|
||||
x_lens,
|
||||
t_tensor,
|
||||
d_tensor,
|
||||
mu,
|
||||
use_grad_ckpt=False,
|
||||
drop_audio_cond=False,
|
||||
drop_text=False,
|
||||
infer=True,
|
||||
text_cache=text_cache,
|
||||
dt_cache=dt_cache,
|
||||
)
|
||||
v_pred = v_pred.transpose(2, 1)
|
||||
if self.use_conditioner_cache:
|
||||
@@ -1037,18 +1149,18 @@ class CFM(torch.nn.Module):
|
||||
dt_cache = dt
|
||||
if inference_cfg_rate > 1e-5:
|
||||
neg, text_cfg_emb, _ = self.estimator(
|
||||
x,
|
||||
prompt_x,
|
||||
x_lens,
|
||||
t_tensor,
|
||||
d_tensor,
|
||||
mu,
|
||||
use_grad_ckpt=False,
|
||||
drop_audio_cond=True,
|
||||
drop_text=True,
|
||||
infer=True,
|
||||
text_cache=text_cfg_cache,
|
||||
dt_cache=dt_cache
|
||||
x,
|
||||
prompt_x,
|
||||
x_lens,
|
||||
t_tensor,
|
||||
d_tensor,
|
||||
mu,
|
||||
use_grad_ckpt=False,
|
||||
drop_audio_cond=True,
|
||||
drop_text=True,
|
||||
infer=True,
|
||||
text_cache=text_cfg_cache,
|
||||
dt_cache=dt_cache,
|
||||
)
|
||||
neg = neg.transpose(2, 1)
|
||||
if self.use_conditioner_cache:
|
||||
@@ -1201,7 +1313,7 @@ class SynthesizerTrnV3(nn.Module):
|
||||
ssl = self.ssl_proj(ssl)
|
||||
quantized, codes, commit_loss, quantized_list = self.quantizer(ssl, layers=[0])
|
||||
quantized = F.interpolate(quantized, scale_factor=2, mode="nearest") ##BCT
|
||||
x, m_p, logs_p, y_mask = self.enc_p(quantized, y_lengths, text, text_lengths, ge)
|
||||
x, m_p, logs_p, y_mask, y_, y_mask_ = self.enc_p(quantized, y_lengths, text, text_lengths, ge)
|
||||
fea = self.bridge(x)
|
||||
fea = F.interpolate(fea, scale_factor=(1.875 if self.version == "v3" else 2), mode="nearest") ##BCT
|
||||
fea, y_mask_ = self.wns1(
|
||||
@@ -1235,7 +1347,7 @@ class SynthesizerTrnV3(nn.Module):
|
||||
quantized = self.quantizer.decode(codes)
|
||||
if self.semantic_frame_rate == "25hz":
|
||||
quantized = F.interpolate(quantized, scale_factor=2, mode="nearest") ##BCT
|
||||
x, m_p, logs_p, y_mask = self.enc_p(quantized, y_lengths, text, text_lengths, ge, speed)
|
||||
x, m_p, logs_p, y_mask, _, _ = self.enc_p(quantized, y_lengths, text, text_lengths, ge, speed)
|
||||
fea = self.bridge(x)
|
||||
fea = F.interpolate(fea, scale_factor=(1.875 if self.version == "v3" else 2), mode="nearest") ##BCT
|
||||
####more wn paramter to learn mel
|
||||
@@ -1352,7 +1464,7 @@ class SynthesizerTrnV3b(nn.Module):
|
||||
ssl = self.ssl_proj(ssl)
|
||||
quantized, codes, commit_loss, quantized_list = self.quantizer(ssl, layers=[0])
|
||||
quantized = F.interpolate(quantized, scale_factor=2, mode="nearest") ##BCT
|
||||
x, m_p, logs_p, y_mask = self.enc_p(quantized, y_lengths, text, text_lengths, ge)
|
||||
x, m_p, logs_p, y_mask, y_, y_mask_ = self.enc_p(quantized, y_lengths, text, text_lengths, ge)
|
||||
z, m_q, logs_q, y_mask = self.enc_q(y, y_lengths, g=ge)
|
||||
z_p = self.flow(z, y_mask, g=ge)
|
||||
z_slice, ids_slice = commons.rand_slice_segments(z, y_lengths, self.segment_size)
|
||||
@@ -1395,7 +1507,7 @@ class SynthesizerTrnV3b(nn.Module):
|
||||
quantized = self.quantizer.decode(codes)
|
||||
if self.semantic_frame_rate == "25hz":
|
||||
quantized = F.interpolate(quantized, scale_factor=2, mode="nearest") ##BCT
|
||||
x, m_p, logs_p, y_mask = self.enc_p(quantized, y_lengths, text, text_lengths, ge)
|
||||
x, m_p, logs_p, y_mask, y_, y_mask_ = self.enc_p(quantized, y_lengths, text, text_lengths, ge)
|
||||
fea = self.bridge(x)
|
||||
fea = F.interpolate(fea, scale_factor=1.875, mode="nearest") ##BCT
|
||||
####more wn paramter to learn mel
|
||||
|
||||
@@ -763,6 +763,9 @@ class CodePredictor(nn.Module):
|
||||
return pred_codes.transpose(0, 1)
|
||||
|
||||
|
||||
v2pro_set = {"v2Pro", "v2ProPlus"}
|
||||
|
||||
|
||||
class SynthesizerTrn(nn.Module):
|
||||
"""
|
||||
Synthesizer for Training
|
||||
@@ -867,20 +870,33 @@ class SynthesizerTrn(nn.Module):
|
||||
# self.enc_p.text_embedding.requires_grad_(False)
|
||||
# self.enc_p.encoder_text.requires_grad_(False)
|
||||
# self.enc_p.mrte.requires_grad_(False)
|
||||
self.is_v2pro = self.version in v2pro_set
|
||||
if self.is_v2pro:
|
||||
self.sv_emb = nn.Linear(20480, gin_channels)
|
||||
self.ge_to512 = nn.Linear(gin_channels, 512)
|
||||
self.prelu = nn.PReLU(num_parameters=gin_channels)
|
||||
|
||||
def forward(self, codes, text, refer, noise_scale=0.5, speed=1):
|
||||
def forward(self, codes, text, refer, noise_scale=0.5, speed=1, sv_emb=None):
|
||||
refer_mask = torch.ones_like(refer[:1, :1, :])
|
||||
if self.version == "v1":
|
||||
ge = self.ref_enc(refer * refer_mask, refer_mask)
|
||||
else:
|
||||
ge = self.ref_enc(refer[:, :704] * refer_mask, refer_mask)
|
||||
if self.is_v2pro:
|
||||
sv_emb = self.sv_emb(sv_emb)
|
||||
ge += sv_emb.unsqueeze(-1)
|
||||
ge = self.prelu(ge)
|
||||
|
||||
quantized = self.quantizer.decode(codes)
|
||||
if self.semantic_frame_rate == "25hz":
|
||||
dquantized = torch.cat([quantized, quantized]).permute(1, 2, 0)
|
||||
quantized = dquantized.contiguous().view(1, self.ssl_dim, -1)
|
||||
|
||||
x, m_p, logs_p, y_mask = self.enc_p(quantized, text, ge, speed)
|
||||
if self.is_v2pro:
|
||||
ge_ = self.ge_to512(ge.transpose(2, 1)).transpose(2, 1)
|
||||
x, m_p, logs_p, y_mask = self.enc_p(quantized, text, ge_, speed)
|
||||
else:
|
||||
x, m_p, logs_p, y_mask = self.enc_p(quantized, text, ge, speed)
|
||||
|
||||
z_p = m_p + torch.randn_like(m_p) * torch.exp(logs_p) * noise_scale
|
||||
|
||||
|
||||
@@ -1,4 +1,5 @@
|
||||
import math
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
from torch import nn
|
||||
@@ -718,8 +719,10 @@ class MelStyleEncoder(nn.Module):
|
||||
else:
|
||||
len_ = (~mask).sum(dim=1).unsqueeze(1)
|
||||
x = x.masked_fill(mask.unsqueeze(-1), 0)
|
||||
x = x.sum(dim=1)
|
||||
out = torch.div(x, len_)
|
||||
dtype = x.dtype
|
||||
x = x.float()
|
||||
x = torch.div(x, len_.unsqueeze(1))
|
||||
out = x.sum(dim=1).to(dtype)
|
||||
return out
|
||||
|
||||
def forward(self, x, mask=None):
|
||||
@@ -743,7 +746,6 @@ class MelStyleEncoder(nn.Module):
|
||||
x = self.fc(x)
|
||||
# temoral average pooling
|
||||
w = self.temporal_avg_pool(x, mask=mask)
|
||||
|
||||
return w.unsqueeze(-1)
|
||||
|
||||
|
||||
|
||||
@@ -10,7 +10,6 @@ i_part = os.environ.get("i_part")
|
||||
all_parts = os.environ.get("all_parts")
|
||||
if "_CUDA_VISIBLE_DEVICES" in os.environ:
|
||||
os.environ["CUDA_VISIBLE_DEVICES"] = os.environ["_CUDA_VISIBLE_DEVICES"]
|
||||
from feature_extractor import cnhubert
|
||||
|
||||
opt_dir = os.environ.get("opt_dir")
|
||||
sv_path = os.environ.get("sv_path")
|
||||
@@ -19,19 +18,18 @@ import torch
|
||||
is_half = eval(os.environ.get("is_half", "True")) and torch.cuda.is_available()
|
||||
|
||||
import traceback
|
||||
import numpy as np
|
||||
from scipy.io import wavfile
|
||||
import torchaudio
|
||||
|
||||
now_dir = os.getcwd()
|
||||
sys.path.append(now_dir)
|
||||
sys.path.append(f"{now_dir}/GPT_SoVITS/eres2net")
|
||||
from tools.my_utils import load_audio, clean_path
|
||||
from tools.my_utils import clean_path
|
||||
from time import time as ttime
|
||||
import shutil
|
||||
from ERes2NetV2 import ERes2NetV2
|
||||
import kaldi as Kaldi
|
||||
|
||||
|
||||
def my_save(fea, path): #####fix issue: torch.save doesn't support chinese path
|
||||
dir = os.path.dirname(path)
|
||||
name = os.path.basename(path)
|
||||
@@ -56,37 +54,45 @@ if torch.cuda.is_available():
|
||||
else:
|
||||
device = "cpu"
|
||||
|
||||
|
||||
class SV:
|
||||
def __init__(self,device,is_half):
|
||||
pretrained_state = torch.load(sv_path, map_location='cpu')
|
||||
embedding_model = ERes2NetV2(baseWidth=24,scale=4,expansion=4)
|
||||
def __init__(self, device, is_half):
|
||||
pretrained_state = torch.load(sv_path, map_location="cpu")
|
||||
embedding_model = ERes2NetV2(baseWidth=24, scale=4, expansion=4)
|
||||
embedding_model.load_state_dict(pretrained_state)
|
||||
embedding_model.eval()
|
||||
self.embedding_model=embedding_model
|
||||
self.res=torchaudio.transforms.Resample(32000, 16000).to(device)
|
||||
self.embedding_model = embedding_model
|
||||
self.res = torchaudio.transforms.Resample(32000, 16000).to(device)
|
||||
if is_half == False:
|
||||
self.embedding_model=self.embedding_model.to(device)
|
||||
self.embedding_model = self.embedding_model.to(device)
|
||||
else:
|
||||
self.embedding_model=self.embedding_model.half().to(device)
|
||||
self.is_half=is_half
|
||||
self.embedding_model = self.embedding_model.half().to(device)
|
||||
self.is_half = is_half
|
||||
|
||||
def compute_embedding3(self,wav):#(1,x)#-1~1
|
||||
def compute_embedding3(self, wav): # (1,x)#-1~1
|
||||
with torch.no_grad():
|
||||
wav=self.res(wav)
|
||||
if self.is_half==True:wav=wav.half()
|
||||
feat = torch.stack([Kaldi.fbank(wav0.unsqueeze(0), num_mel_bins=80, sample_frequency=16000, dither=0) for wav0 in wav])
|
||||
wav = self.res(wav)
|
||||
if self.is_half == True:
|
||||
wav = wav.half()
|
||||
feat = torch.stack(
|
||||
[Kaldi.fbank(wav0.unsqueeze(0), num_mel_bins=80, sample_frequency=16000, dither=0) for wav0 in wav]
|
||||
)
|
||||
sv_emb = self.embedding_model.forward3(feat)
|
||||
return sv_emb
|
||||
|
||||
sv=SV(device,is_half)
|
||||
|
||||
sv = SV(device, is_half)
|
||||
|
||||
|
||||
def name2go(wav_name, wav_path):
|
||||
sv_cn_path = "%s/%s.pt" % (sv_cn_dir, wav_name)
|
||||
if os.path.exists(sv_cn_path):return
|
||||
wav_path="%s/%s" % (wav32dir, wav_name)
|
||||
wav32k,sr0 = torchaudio.load(wav_path)
|
||||
assert sr0==32000
|
||||
if os.path.exists(sv_cn_path):
|
||||
return
|
||||
wav_path = "%s/%s" % (wav32dir, wav_name)
|
||||
wav32k, sr0 = torchaudio.load(wav_path)
|
||||
assert sr0 == 32000
|
||||
wav32k = wav32k.to(device)
|
||||
emb=sv.compute_embedding3(wav32k).cpu() # torch.Size([1, 20480])
|
||||
emb = sv.compute_embedding3(wav32k).cpu() # torch.Size([1, 20480])
|
||||
my_save(emb, sv_cn_path)
|
||||
|
||||
|
||||
|
||||
@@ -17,15 +17,16 @@ def my_save(fea, path): #####fix issue: torch.save doesn't support chinese path
|
||||
shutil.move(tmp_path, "%s/%s" % (dir, name))
|
||||
|
||||
|
||||
|
||||
from io import BytesIO
|
||||
|
||||
model_version2byte={
|
||||
"v3":b"03",
|
||||
"v4":b"04",
|
||||
"v2Pro":b"05",
|
||||
"v2ProPlus":b"06",
|
||||
model_version2byte = {
|
||||
"v3": b"03",
|
||||
"v4": b"04",
|
||||
"v2Pro": b"05",
|
||||
"v2ProPlus": b"06",
|
||||
}
|
||||
|
||||
|
||||
def my_save2(fea, path, model_version):
|
||||
bio = BytesIO()
|
||||
torch.save(fea, bio)
|
||||
@@ -50,7 +51,7 @@ def savee(ckpt, name, epoch, steps, hps, model_version=None, lora_rank=None):
|
||||
if lora_rank:
|
||||
opt["lora_rank"] = lora_rank
|
||||
my_save2(opt, "%s/%s.pth" % (hps.save_weight_dir, name), model_version)
|
||||
elif (model_version!=None and "Pro"in model_version):
|
||||
elif model_version != None and "Pro" in model_version:
|
||||
my_save2(opt, "%s/%s.pth" % (hps.save_weight_dir, name), model_version)
|
||||
else:
|
||||
my_save(opt, "%s/%s.pth" % (hps.save_weight_dir, name))
|
||||
@@ -58,6 +59,7 @@ def savee(ckpt, name, epoch, steps, hps, model_version=None, lora_rank=None):
|
||||
except:
|
||||
return traceback.format_exc()
|
||||
|
||||
|
||||
"""
|
||||
00:v1
|
||||
01:v2
|
||||
@@ -127,7 +129,7 @@ def get_sovits_version_from_path_fast(sovits_path):
|
||||
def load_sovits_new(sovits_path):
|
||||
f = open(sovits_path, "rb")
|
||||
meta = f.read(2)
|
||||
if meta != "PK":
|
||||
if meta != b"PK":
|
||||
data = b"PK" + f.read()
|
||||
bio = BytesIO()
|
||||
bio.write(data)
|
||||
|
||||
+79
-16
@@ -36,7 +36,7 @@ from module.models import (
|
||||
MultiPeriodDiscriminator,
|
||||
SynthesizerTrn,
|
||||
)
|
||||
from process_ckpt import savee,my_save2
|
||||
from process_ckpt import savee
|
||||
|
||||
torch.backends.cudnn.benchmark = False
|
||||
torch.backends.cudnn.deterministic = False
|
||||
@@ -87,11 +87,30 @@ def run(rank, n_gpus, hps):
|
||||
if torch.cuda.is_available():
|
||||
torch.cuda.set_device(rank)
|
||||
|
||||
train_dataset = TextAudioSpeakerLoader(hps.data,version=hps.model.version)
|
||||
train_dataset = TextAudioSpeakerLoader(hps.data, version=hps.model.version)
|
||||
train_sampler = DistributedBucketSampler(
|
||||
train_dataset,
|
||||
hps.train.batch_size,
|
||||
[32,300,400,500,600,700,800,900,1000,1100,1200,1300,1400,1500,1600,1700,1800,1900,],
|
||||
[
|
||||
32,
|
||||
300,
|
||||
400,
|
||||
500,
|
||||
600,
|
||||
700,
|
||||
800,
|
||||
900,
|
||||
1000,
|
||||
1100,
|
||||
1200,
|
||||
1300,
|
||||
1400,
|
||||
1500,
|
||||
1600,
|
||||
1700,
|
||||
1800,
|
||||
1900,
|
||||
],
|
||||
num_replicas=n_gpus,
|
||||
rank=rank,
|
||||
shuffle=True,
|
||||
@@ -105,7 +124,7 @@ def run(rank, n_gpus, hps):
|
||||
collate_fn=collate_fn,
|
||||
batch_sampler=train_sampler,
|
||||
persistent_workers=True,
|
||||
prefetch_factor=4,
|
||||
prefetch_factor=3,
|
||||
)
|
||||
# if rank == 0:
|
||||
# eval_dataset = TextAudioSpeakerLoader(hps.data.validation_files, hps.data, val=True)
|
||||
@@ -130,9 +149,9 @@ def run(rank, n_gpus, hps):
|
||||
)
|
||||
|
||||
net_d = (
|
||||
MultiPeriodDiscriminator(hps.model.use_spectral_norm,version=hps.model.version).cuda(rank)
|
||||
MultiPeriodDiscriminator(hps.model.use_spectral_norm, version=hps.model.version).cuda(rank)
|
||||
if torch.cuda.is_available()
|
||||
else MultiPeriodDiscriminator(hps.model.use_spectral_norm,version=hps.model.version).to(device)
|
||||
else MultiPeriodDiscriminator(hps.model.use_spectral_norm, version=hps.model.version).to(device)
|
||||
)
|
||||
for name, param in net_g.named_parameters():
|
||||
if not param.requires_grad:
|
||||
@@ -235,7 +254,7 @@ def run(rank, n_gpus, hps):
|
||||
print(
|
||||
"loaded pretrained %s" % hps.train.pretrained_s2D,
|
||||
net_d.module.load_state_dict(
|
||||
torch.load(hps.train.pretrained_s2D, map_location="cpu", weights_only=False)["weight"],strict=False
|
||||
torch.load(hps.train.pretrained_s2D, map_location="cpu", weights_only=False)["weight"], strict=False
|
||||
)
|
||||
if torch.cuda.is_available()
|
||||
else net_d.load_state_dict(
|
||||
@@ -310,17 +329,44 @@ def train_and_evaluate(rank, epoch, hps, nets, optims, schedulers, scaler, loade
|
||||
net_g.train()
|
||||
net_d.train()
|
||||
for batch_idx, data in enumerate(tqdm(train_loader)):
|
||||
if hps.model.version in {"v2Pro","v2ProPlus"}:
|
||||
ssl, ssl_lengths, spec, spec_lengths, y, y_lengths, text, text_lengths,sv_emb=data
|
||||
if hps.model.version in {"v2Pro", "v2ProPlus"}:
|
||||
ssl, ssl_lengths, spec, spec_lengths, y, y_lengths, text, text_lengths, sv_emb = data
|
||||
else:
|
||||
ssl, ssl_lengths, spec, spec_lengths, y, y_lengths, text, text_lengths=data
|
||||
ssl, ssl_lengths, spec, spec_lengths, y, y_lengths, text, text_lengths = data
|
||||
if torch.cuda.is_available():
|
||||
spec, spec_lengths = (spec.cuda(rank,non_blocking=True,),spec_lengths.cuda(rank,non_blocking=True,),)
|
||||
y, y_lengths = (y.cuda(rank,non_blocking=True,),y_lengths.cuda(rank,non_blocking=True,),)
|
||||
spec, spec_lengths = (
|
||||
spec.cuda(
|
||||
rank,
|
||||
non_blocking=True,
|
||||
),
|
||||
spec_lengths.cuda(
|
||||
rank,
|
||||
non_blocking=True,
|
||||
),
|
||||
)
|
||||
y, y_lengths = (
|
||||
y.cuda(
|
||||
rank,
|
||||
non_blocking=True,
|
||||
),
|
||||
y_lengths.cuda(
|
||||
rank,
|
||||
non_blocking=True,
|
||||
),
|
||||
)
|
||||
ssl = ssl.cuda(rank, non_blocking=True)
|
||||
ssl.requires_grad = False
|
||||
# ssl_lengths = ssl_lengths.cuda(rank, non_blocking=True)
|
||||
text, text_lengths = (text.cuda(rank,non_blocking=True,),text_lengths.cuda(rank,non_blocking=True,),)
|
||||
text, text_lengths = (
|
||||
text.cuda(
|
||||
rank,
|
||||
non_blocking=True,
|
||||
),
|
||||
text_lengths.cuda(
|
||||
rank,
|
||||
non_blocking=True,
|
||||
),
|
||||
)
|
||||
if hps.model.version in {"v2Pro", "v2ProPlus"}:
|
||||
sv_emb = sv_emb.cuda(rank, non_blocking=True)
|
||||
else:
|
||||
@@ -334,9 +380,19 @@ def train_and_evaluate(rank, epoch, hps, nets, optims, schedulers, scaler, loade
|
||||
sv_emb = sv_emb.to(device)
|
||||
with autocast(enabled=hps.train.fp16_run):
|
||||
if hps.model.version in {"v2Pro", "v2ProPlus"}:
|
||||
(y_hat,kl_ssl,ids_slice,x_mask,z_mask,(z, z_p, m_p, logs_p, m_q, logs_q),stats_ssl) = net_g(ssl, spec, spec_lengths, text, text_lengths,sv_emb)
|
||||
(y_hat, kl_ssl, ids_slice, x_mask, z_mask, (z, z_p, m_p, logs_p, m_q, logs_q), stats_ssl) = net_g(
|
||||
ssl, spec, spec_lengths, text, text_lengths, sv_emb
|
||||
)
|
||||
else:
|
||||
(y_hat,kl_ssl,ids_slice,x_mask,z_mask,(z, z_p, m_p, logs_p, m_q, logs_q),stats_ssl,) = net_g(ssl, spec, spec_lengths, text, text_lengths)
|
||||
(
|
||||
y_hat,
|
||||
kl_ssl,
|
||||
ids_slice,
|
||||
x_mask,
|
||||
z_mask,
|
||||
(z, z_p, m_p, logs_p, m_q, logs_q),
|
||||
stats_ssl,
|
||||
) = net_g(ssl, spec, spec_lengths, text, text_lengths)
|
||||
|
||||
mel = spec_to_mel_torch(
|
||||
spec,
|
||||
@@ -508,7 +564,14 @@ def train_and_evaluate(rank, epoch, hps, nets, optims, schedulers, scaler, loade
|
||||
% (
|
||||
hps.name,
|
||||
epoch,
|
||||
savee(ckpt,hps.name + "_e%s_s%s" % (epoch, global_step),epoch,global_step,hps,model_version=None if hps.model.version not in {"v2Pro","v2ProPlus"}else hps.model.version),
|
||||
savee(
|
||||
ckpt,
|
||||
hps.name + "_e%s_s%s" % (epoch, global_step),
|
||||
epoch,
|
||||
global_step,
|
||||
hps,
|
||||
model_version=None if hps.model.version not in {"v2Pro", "v2ProPlus"} else hps.model.version,
|
||||
),
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
@@ -118,13 +118,13 @@ def run(rank, n_gpus, hps):
|
||||
collate_fn = TextAudioSpeakerCollate()
|
||||
train_loader = DataLoader(
|
||||
train_dataset,
|
||||
num_workers=6,
|
||||
num_workers=5,
|
||||
shuffle=False,
|
||||
pin_memory=True,
|
||||
collate_fn=collate_fn,
|
||||
batch_sampler=train_sampler,
|
||||
persistent_workers=True,
|
||||
prefetch_factor=4,
|
||||
prefetch_factor=3,
|
||||
)
|
||||
# if rank == 0:
|
||||
# eval_dataset = TextAudioSpeakerLoader(hps.data.validation_files, hps.data, val=True)
|
||||
|
||||
@@ -55,6 +55,10 @@ def main():
|
||||
n_gpus = torch.cuda.device_count()
|
||||
else:
|
||||
n_gpus = 1
|
||||
if n_gpus <= 1:
|
||||
run(0, n_gpus, hps)
|
||||
return
|
||||
|
||||
os.environ["MASTER_ADDR"] = "localhost"
|
||||
os.environ["MASTER_PORT"] = str(randint(20000, 55555))
|
||||
|
||||
@@ -77,12 +81,14 @@ def run(rank, n_gpus, hps):
|
||||
writer = SummaryWriter(log_dir=hps.s2_ckpt_dir)
|
||||
writer_eval = SummaryWriter(log_dir=os.path.join(hps.s2_ckpt_dir, "eval"))
|
||||
|
||||
dist.init_process_group(
|
||||
backend="gloo" if os.name == "nt" or not torch.cuda.is_available() else "nccl",
|
||||
init_method="env://?use_libuv=False",
|
||||
world_size=n_gpus,
|
||||
rank=rank,
|
||||
)
|
||||
use_ddp = n_gpus > 1
|
||||
if use_ddp:
|
||||
dist.init_process_group(
|
||||
backend="gloo" if os.name == "nt" or not torch.cuda.is_available() else "nccl",
|
||||
init_method="env://?use_libuv=False",
|
||||
world_size=n_gpus,
|
||||
rank=rank,
|
||||
)
|
||||
torch.manual_seed(hps.train.seed)
|
||||
if torch.cuda.is_available():
|
||||
torch.cuda.set_device(rank)
|
||||
@@ -118,15 +124,20 @@ def run(rank, n_gpus, hps):
|
||||
shuffle=True,
|
||||
)
|
||||
collate_fn = TextAudioSpeakerCollate()
|
||||
train_loader = DataLoader(
|
||||
train_dataset,
|
||||
num_workers=6,
|
||||
worker_count = 0 if os.name == "nt" and n_gpus <= 1 else min(2 if os.name == "nt" else 5, os.cpu_count() or 1)
|
||||
loader_kwargs = dict(
|
||||
num_workers=worker_count,
|
||||
shuffle=False,
|
||||
pin_memory=True,
|
||||
pin_memory=torch.cuda.is_available(),
|
||||
collate_fn=collate_fn,
|
||||
batch_sampler=train_sampler,
|
||||
persistent_workers=True,
|
||||
prefetch_factor=4,
|
||||
)
|
||||
if worker_count > 0:
|
||||
loader_kwargs["persistent_workers"] = True
|
||||
loader_kwargs["prefetch_factor"] = 2 if os.name == "nt" else 3
|
||||
train_loader = DataLoader(
|
||||
train_dataset,
|
||||
**loader_kwargs,
|
||||
)
|
||||
save_root = "%s/logs_s2_%s_lora_%s" % (hps.data.exp_dir, hps.model.version, hps.train.lora_rank)
|
||||
os.makedirs(save_root, exist_ok=True)
|
||||
@@ -156,7 +167,9 @@ def run(rank, n_gpus, hps):
|
||||
|
||||
def model2cuda(net_g, rank):
|
||||
if torch.cuda.is_available():
|
||||
net_g = DDP(net_g.cuda(rank), device_ids=[rank], find_unused_parameters=True)
|
||||
net_g = net_g.cuda(rank)
|
||||
if use_ddp:
|
||||
net_g = DDP(net_g, device_ids=[rank], find_unused_parameters=True)
|
||||
else:
|
||||
net_g = net_g.to(device)
|
||||
return net_g
|
||||
@@ -242,6 +255,8 @@ def run(rank, n_gpus, hps):
|
||||
None,
|
||||
)
|
||||
scheduler_g.step()
|
||||
if use_ddp and dist.is_initialized():
|
||||
dist.destroy_process_group()
|
||||
print("training done")
|
||||
|
||||
|
||||
|
||||
@@ -0,0 +1,611 @@
|
||||
# 这是一个实验性质的实现,旨在探索 stream infer 的可能性。(xiao hai xie zhe wan de)
|
||||
from typing import List
|
||||
from export_torch_script import ExportERes2NetV2, SSLModel, T2SModel, VitsModel, get_raw_t2s_model, init_sv_cn, resamplex, sample, spectrogram_torch
|
||||
import export_torch_script
|
||||
from my_utils import load_audio
|
||||
import torch
|
||||
from torch import LongTensor, Tensor, nn
|
||||
from torch.nn import functional as F
|
||||
|
||||
import soundfile
|
||||
from inference_webui import get_phones_and_bert
|
||||
import matplotlib.pyplot as plt
|
||||
|
||||
|
||||
class StreamT2SModel(nn.Module):
|
||||
def __init__(self, t2s: T2SModel):
|
||||
super(StreamT2SModel, self).__init__()
|
||||
self.t2s = t2s
|
||||
|
||||
@torch.jit.export
|
||||
def pre_infer(
|
||||
self,
|
||||
prompts: LongTensor,
|
||||
ref_seq: LongTensor,
|
||||
text_seq: LongTensor,
|
||||
ref_bert: torch.Tensor,
|
||||
text_bert: torch.Tensor,
|
||||
top_k: int,
|
||||
) -> tuple[int, Tensor, Tensor, List[Tensor], List[Tensor]]:
|
||||
bert = torch.cat([ref_bert.T, text_bert.T], 1)
|
||||
all_phoneme_ids = torch.cat([ref_seq, text_seq], 1)
|
||||
bert = bert.unsqueeze(0)
|
||||
|
||||
x = self.t2s.ar_text_embedding(all_phoneme_ids)
|
||||
x = x + self.t2s.bert_proj(bert.transpose(1, 2))
|
||||
x: torch.Tensor = self.t2s.ar_text_position(x)
|
||||
|
||||
# [1,N,512] [1,N]
|
||||
# y, k, v, y_emb, x_example = self.first_stage_decoder(x, prompts)
|
||||
y = prompts
|
||||
# x_example = x[:,:,0] * 0.0
|
||||
|
||||
x_len = x.shape[1]
|
||||
x_attn_mask = torch.zeros((x_len, x_len), dtype=torch.bool)
|
||||
|
||||
y_emb = self.t2s.ar_audio_embedding(y)
|
||||
y_len: int = y_emb.shape[1]
|
||||
prefix_len = y.shape[1]
|
||||
y_pos = self.t2s.ar_audio_position(y_emb)
|
||||
xy_pos = torch.concat([x, y_pos], dim=1)
|
||||
|
||||
bsz = x.shape[0]
|
||||
src_len = x_len + y_len
|
||||
x_attn_mask_pad = F.pad(
|
||||
x_attn_mask,
|
||||
(0, y_len), ###xx的纯0扩展到xx纯0+xy纯1,(x,x+y)
|
||||
value=True,
|
||||
)
|
||||
y_attn_mask = F.pad( ###yy的右上1扩展到左边xy的0,(y,x+y)
|
||||
torch.triu(torch.ones(y_len, y_len, dtype=torch.bool), diagonal=1),
|
||||
(x_len, 0),
|
||||
value=False,
|
||||
)
|
||||
xy_attn_mask = (
|
||||
torch.concat([x_attn_mask_pad, y_attn_mask], dim=0)
|
||||
.unsqueeze(0)
|
||||
.expand(bsz * self.t2s.num_head, -1, -1)
|
||||
.view(bsz, self.t2s.num_head, src_len, src_len)
|
||||
.to(device=x.device, dtype=torch.bool)
|
||||
)
|
||||
|
||||
xy_dec, k_cache, v_cache = self.t2s.t2s_transformer.process_prompt(
|
||||
xy_pos, xy_attn_mask, None
|
||||
)
|
||||
|
||||
logits = self.t2s.ar_predict_layer(xy_dec[:, -1])
|
||||
logits = logits[:, :-1]
|
||||
samples = sample(
|
||||
logits, y, top_k=top_k, top_p=1, repetition_penalty=1.35, temperature=1.0
|
||||
)[0]
|
||||
y = torch.concat([y, samples], dim=1)
|
||||
y_emb: Tensor = self.t2s.ar_audio_embedding(y[:, -1:])
|
||||
xy_pos: Tensor = (
|
||||
y_emb * self.t2s.ar_audio_position.x_scale
|
||||
+ self.t2s.ar_audio_position.alpha
|
||||
* self.t2s.ar_audio_position.pe[:, y_len].to(
|
||||
dtype=y_emb.dtype, device=y_emb.device
|
||||
)
|
||||
)
|
||||
|
||||
return y_len, y, xy_pos, k_cache, v_cache
|
||||
|
||||
@torch.jit.export
|
||||
def decode_next_token(
|
||||
self,
|
||||
idx: int, # 记住从1开始 到1500
|
||||
top_k: int,
|
||||
y_len: int,
|
||||
y: Tensor,
|
||||
xy_pos: Tensor,
|
||||
k_cache: List[Tensor],
|
||||
v_cache: List[Tensor],
|
||||
) -> tuple[Tensor, Tensor, int, List[Tensor], List[Tensor]]:
|
||||
# [1, N] [N_layer, N, 1, 512] [N_layer, N, 1, 512] [1, N, 512] [1] [1, N, 512] [1, N]
|
||||
# y, k, v, y_emb, logits, samples = self.stage_decoder(y, k, v, y_emb, x_example)
|
||||
xy_dec, k_cache, v_cache = self.t2s.t2s_transformer.decode_next_token(
|
||||
xy_pos, k_cache, v_cache
|
||||
)
|
||||
logits = self.t2s.ar_predict_layer(xy_dec[:, -1])
|
||||
|
||||
if idx < 11: ###至少预测出10个token不然不给停止(0.4s)
|
||||
logits = logits[:, :-1]
|
||||
|
||||
samples = sample(
|
||||
logits, y, top_k=top_k, top_p=1, repetition_penalty=1.35, temperature=1.0
|
||||
)[0]
|
||||
|
||||
y = torch.concat([y, samples], dim=1)
|
||||
last_token = int(samples[0, 0])
|
||||
|
||||
# if early_stop_num != -1 and (y.shape[1] - prefix_len) > early_stop_num:
|
||||
# stop = True
|
||||
if torch.argmax(logits, dim=-1)[0] == self.t2s.EOS or samples[0, 0] == self.t2s.EOS:
|
||||
return y[:,:-1], xy_pos, self.t2s.EOS, k_cache, v_cache
|
||||
|
||||
# if stop:
|
||||
# if y.shape[1] == 0:
|
||||
# y = torch.concat([y, torch.zeros_like(samples)], dim=1)
|
||||
# break
|
||||
|
||||
y_emb = self.t2s.ar_audio_embedding(y[:, -1:])
|
||||
xy_pos = (
|
||||
y_emb * self.t2s.ar_audio_position.x_scale
|
||||
+ self.t2s.ar_audio_position.alpha
|
||||
* self.t2s.ar_audio_position.pe[:, y_len + idx].to(
|
||||
dtype=y_emb.dtype, device=y_emb.device
|
||||
)
|
||||
)
|
||||
return y, xy_pos, last_token, k_cache, v_cache
|
||||
|
||||
def forward(
|
||||
self,
|
||||
idx: int, # 记住从1开始 到1500
|
||||
top_k: int,
|
||||
y_len: int,
|
||||
y: Tensor,
|
||||
xy_pos: Tensor,
|
||||
k_cache: List[Tensor],
|
||||
v_cache: List[Tensor],
|
||||
):
|
||||
return self.decode_next_token(idx,top_k,y_len,y,xy_pos,k_cache,v_cache)
|
||||
|
||||
|
||||
class StepVitsModel(nn.Module):
|
||||
def __init__(self, vits: VitsModel,sv_model:ExportERes2NetV2):
|
||||
super().__init__()
|
||||
self.hps = vits.hps
|
||||
self.vq_model = vits.vq_model
|
||||
self.hann_window = vits.hann_window
|
||||
self.sv = sv_model
|
||||
|
||||
def ref_handle(self, ref_audio_32k):
|
||||
refer = spectrogram_torch(
|
||||
self.hann_window,
|
||||
ref_audio_32k.float(),
|
||||
self.hps.data.filter_length,
|
||||
self.hps.data.sampling_rate,
|
||||
self.hps.data.hop_length,
|
||||
self.hps.data.win_length,
|
||||
center=False,
|
||||
)
|
||||
refer = refer.to(ref_audio_32k.dtype)
|
||||
ref_audio_16k = resamplex(ref_audio_32k, 32000, 16000).to(ref_audio_32k.dtype).to(ref_audio_32k.device)
|
||||
sv_emb = self.sv(ref_audio_16k)
|
||||
return refer, sv_emb
|
||||
|
||||
def extract_latent(self, ssl_content):
|
||||
codes = self.vq_model.extract_latent(ssl_content)
|
||||
return codes[0]
|
||||
|
||||
def forward(self, pred_semantic, text_seq, refer, sv_emb=None):
|
||||
return self.vq_model(
|
||||
pred_semantic, text_seq, refer, speed=1.0, sv_emb=sv_emb
|
||||
)[0, 0]
|
||||
|
||||
|
||||
@torch.jit.script
|
||||
def find_best_audio_offset_fast(reference_audio: Tensor, search_audio: Tensor):
|
||||
ref_len = len(reference_audio)
|
||||
search_len = len(search_audio)
|
||||
|
||||
if search_len < ref_len:
|
||||
raise ValueError(
|
||||
f"搜索音频长度 ({search_len}) 必须大于等于参考音频长度 ({ref_len})"
|
||||
)
|
||||
|
||||
# 使用F.conv1d计算原始互相关
|
||||
reference_flipped = reference_audio.unsqueeze(0).unsqueeze(0)
|
||||
search_padded = search_audio.unsqueeze(0).unsqueeze(0)
|
||||
|
||||
# 计算点积
|
||||
dot_products = F.conv1d(search_padded, reference_flipped).squeeze()
|
||||
|
||||
if len(dot_products.shape) == 0:
|
||||
dot_products = dot_products.unsqueeze(0)
|
||||
|
||||
# 计算参考音频的平方和
|
||||
ref_squared_sum = torch.sum(reference_audio**2)
|
||||
|
||||
# 计算搜索音频每个位置的平方和(滑动窗口)
|
||||
search_squared = search_audio**2
|
||||
search_squared_padded = search_squared.unsqueeze(0).unsqueeze(0)
|
||||
ones_kernel = torch.ones(
|
||||
1, 1, ref_len, dtype=search_audio.dtype, device=search_audio.device
|
||||
)
|
||||
|
||||
segment_squared_sums = F.conv1d(search_squared_padded, ones_kernel).squeeze()
|
||||
|
||||
if len(segment_squared_sums.shape) == 0:
|
||||
segment_squared_sums = segment_squared_sums.unsqueeze(0)
|
||||
|
||||
# 计算归一化因子
|
||||
ref_norm = torch.sqrt(ref_squared_sum)
|
||||
segment_norms = torch.sqrt(segment_squared_sums)
|
||||
|
||||
# 避免除零
|
||||
epsilon = 1e-8
|
||||
normalization_factor = ref_norm * segment_norms + epsilon
|
||||
|
||||
# 归一化互相关
|
||||
correlation_scores = dot_products / normalization_factor
|
||||
|
||||
best_offset = torch.argmax(correlation_scores).item()
|
||||
|
||||
return best_offset, correlation_scores
|
||||
|
||||
|
||||
import time
|
||||
|
||||
def test_stream(
|
||||
gpt_path,
|
||||
vits_path,
|
||||
version,
|
||||
ref_audio_path,
|
||||
ref_text,
|
||||
output_path,
|
||||
device="cpu",
|
||||
is_half=True,
|
||||
):
|
||||
if export_torch_script.sv_cn_model == None:
|
||||
init_sv_cn(device,is_half)
|
||||
|
||||
ref_audio = torch.tensor([load_audio(ref_audio_path, 16000)]).float()
|
||||
ssl = SSLModel()
|
||||
|
||||
print(f"device: {device}")
|
||||
|
||||
ref_seq_id, ref_bert_T, ref_norm_text = get_phones_and_bert(
|
||||
ref_text, "all_zh", "v2"
|
||||
)
|
||||
ref_seq = torch.LongTensor([ref_seq_id]).to(device)
|
||||
ref_bert = ref_bert_T.T
|
||||
if is_half:
|
||||
ref_bert = ref_bert.half()
|
||||
ref_bert = ref_bert.to(ref_seq.device)
|
||||
|
||||
text_seq_id, text_bert_T, norm_text = get_phones_and_bert(
|
||||
"这是一个简单的示例,真没想到这么简单就完成了,真的神奇,接下来我们说说狐狸,可能这就是狐狸吧.它有长长的尾巴,尖尖的耳朵,传说中还有九条尾巴。你觉得狐狸神奇吗?", "auto", "v2"
|
||||
)
|
||||
text_seq = torch.LongTensor([text_seq_id]).to(device)
|
||||
text_bert = text_bert_T.T
|
||||
if is_half:
|
||||
text_bert = text_bert.half()
|
||||
text_bert = text_bert.to(text_seq.device)
|
||||
|
||||
ssl_content = ssl(ref_audio)
|
||||
if is_half:
|
||||
ssl_content = ssl_content.half()
|
||||
ssl_content = ssl_content.to(device)
|
||||
|
||||
sv_model = ExportERes2NetV2(export_torch_script.sv_cn_model)
|
||||
|
||||
# vits_path = "SoVITS_weights_v2/xw_e8_s216.pth"
|
||||
vits = VitsModel(vits_path, version,is_half=is_half,device=device)
|
||||
vits.eval()
|
||||
|
||||
# gpt_path = "GPT_weights_v2/xw-e15.ckpt"
|
||||
# dict_s1 = torch.load(gpt_path, map_location=device)
|
||||
dict_s1 = torch.load(gpt_path, weights_only=False)
|
||||
raw_t2s = get_raw_t2s_model(dict_s1).to(device)
|
||||
print("#### get_raw_t2s_model ####")
|
||||
print(raw_t2s.config)
|
||||
if is_half:
|
||||
raw_t2s = raw_t2s.half()
|
||||
t2s_m = T2SModel(raw_t2s)
|
||||
t2s_m.eval()
|
||||
# t2s = torch.jit.script(t2s_m).to(device)
|
||||
t2s = t2s_m
|
||||
print("#### script t2s_m ####")
|
||||
|
||||
print("vits.hps.data.sampling_rate:", vits.hps.data.sampling_rate)
|
||||
|
||||
stream_t2s = StreamT2SModel(t2s).to(device)
|
||||
stream_t2s = torch.jit.script(stream_t2s)
|
||||
|
||||
ref_audio_sr = resamplex(ref_audio, 16000, 32000)
|
||||
if is_half:
|
||||
ref_audio_sr = ref_audio_sr.half()
|
||||
ref_audio_sr = ref_audio_sr.to(device)
|
||||
|
||||
top_k = 15
|
||||
|
||||
codes = vits.vq_model.extract_latent(ssl_content)
|
||||
prompt_semantic = codes[0, 0]
|
||||
prompts = prompt_semantic.unsqueeze(0)
|
||||
|
||||
audio_16k = resamplex(ref_audio_sr, 32000, 16000).to(ref_audio_sr.dtype)
|
||||
sv_emb = sv_model(audio_16k)
|
||||
print("text_seq",text_seq.shape)
|
||||
|
||||
refer = spectrogram_torch(
|
||||
vits.hann_window,
|
||||
ref_audio_sr,
|
||||
vits.hps.data.filter_length,
|
||||
vits.hps.data.sampling_rate,
|
||||
vits.hps.data.hop_length,
|
||||
vits.hps.data.win_length,
|
||||
center=False,
|
||||
)
|
||||
|
||||
st = time.time()
|
||||
et = time.time()
|
||||
|
||||
y_len, y, xy_pos, k_cache, v_cache = stream_t2s.pre_infer(prompts, ref_seq, text_seq, ref_bert, text_bert, top_k)
|
||||
idx = 1
|
||||
last_idx = 0
|
||||
audios = []
|
||||
raw_audios = []
|
||||
last_audio_ret = None
|
||||
offset_index = []
|
||||
full_audios = []
|
||||
print("y.shape:", y.shape)
|
||||
cut_id = 0
|
||||
while True:
|
||||
y, xy_pos, last_token, k_cache, v_cache = stream_t2s(idx, top_k, y_len, y, xy_pos, k_cache, v_cache)
|
||||
# print("y.shape:", y.shape)
|
||||
stop = last_token==t2s.EOS
|
||||
print('idx:',idx , 'y.shape:', y.shape, y.shape[1]-idx)
|
||||
|
||||
if last_token < 50 and idx-last_idx > (len(audios)+1) * 25 and idx > cut_id:
|
||||
cut_id = idx + 7
|
||||
print('trigger:',idx, last_idx, y[:,-idx+last_idx:], y[:,-idx+last_idx:].shape)
|
||||
# y = torch.cat([y, y[:,-1:]], dim=1)
|
||||
# idx+=1
|
||||
|
||||
if stop :
|
||||
idx -=1
|
||||
print('stop')
|
||||
print(idx, y[:,-idx+last_idx:])
|
||||
print(idx,last_idx, y.shape)
|
||||
print(y[:,-idx:-idx+20])
|
||||
|
||||
|
||||
# 玄学这档子事说不清楚
|
||||
if idx == cut_id or stop:
|
||||
print(f"idx: {idx}, last_idx: {last_idx}, cut_id: {cut_id}, stop: {stop}")
|
||||
audio = vits.vq_model(y[:,-idx:].unsqueeze(0), text_seq, refer, speed=1.0, sv_emb=sv_emb)[0, 0]
|
||||
full_audios.append(audio)
|
||||
if last_idx == 0:
|
||||
last_audio_ret = audio[-1280*8:-1280*8+256]
|
||||
audio = audio[:-1280*8]
|
||||
raw_audios.append(audio)
|
||||
et = time.time()
|
||||
else:
|
||||
if stop:
|
||||
audio_ = audio[last_idx*1280 -1280*8:]
|
||||
raw_audios.append(audio_)
|
||||
i, x = find_best_audio_offset_fast(last_audio_ret, audio_[:1280])
|
||||
offset_index.append(i)
|
||||
audio = audio_[i:]
|
||||
else:
|
||||
audio_ = audio[last_idx*1280 -1280*8:-1280*8]
|
||||
raw_audios.append(audio_)
|
||||
i, x = find_best_audio_offset_fast(last_audio_ret, audio_[:1280])
|
||||
offset_index.append(i)
|
||||
last_audio_ret = audio[-1280*8:-1280*8+256]
|
||||
audio = audio_[i:]
|
||||
last_idx = idx
|
||||
# print(f'write {output_path}/out_{audio_index}')
|
||||
# soundfile.write(f"{output_path}/out_{audio_index}.wav", audio.float().detach().cpu().numpy(), 32000)
|
||||
audios.append(audio)
|
||||
# print(idx,'/',1500 , y.shape, y[0,-1].item(), stop)
|
||||
if idx>1500:
|
||||
break
|
||||
|
||||
if stop:
|
||||
break
|
||||
|
||||
idx+=1
|
||||
|
||||
at = time.time()
|
||||
|
||||
for (i,a) in enumerate(audios):
|
||||
print(f'write {output_path}/out_{i}')
|
||||
soundfile.write(f"{output_path}/out_{i}.wav", a.float().detach().cpu().numpy(), 32000)
|
||||
|
||||
print(f"frist token: {et - st:.4f} seconds")
|
||||
print(f"all token: {at - st:.4f} seconds")
|
||||
audio = vits.vq_model(y[:,-idx:].unsqueeze(0), text_seq, refer, speed=1.0, sv_emb=sv_emb)[0, 0]
|
||||
soundfile.write(f"{output_path}/out_final.wav", audio.float().detach().cpu().numpy(), 32000)
|
||||
audio = torch.cat(audios, dim=0)
|
||||
soundfile.write(f"{output_path}/out.wav", audio.float().detach().cpu().numpy(), 32000)
|
||||
audio_raw = torch.cat(raw_audios, dim=0)
|
||||
soundfile.write(f"{output_path}/out.raw.wav", audio_raw.float().detach().cpu().numpy(), 32000)
|
||||
|
||||
|
||||
colors = ['red', 'green', 'blue', 'orange', 'purple', 'cyan', 'magenta', 'yellow']
|
||||
|
||||
max_duration = full_audios[-1].shape[0]
|
||||
plt.xlim(0, max_duration)
|
||||
|
||||
last_line = 0
|
||||
|
||||
for i,a in enumerate(full_audios):
|
||||
plt.plot((a+2.0*i).float().detach().cpu().numpy(), color=colors[i], alpha=0.5, label=f"Audio {i}")
|
||||
# plt.axvline(x=last_line, color=colors[i], linestyle='--')
|
||||
last_line = a.shape[0]-8*1280
|
||||
plt.axvline(x=last_line, color=colors[i], linestyle='--')
|
||||
|
||||
plt.plot((audio-2.0).float().detach().cpu().numpy(), color='black', label='Final Audio')
|
||||
|
||||
plt.plot((audio_raw-4.0).float().detach().cpu().numpy(), color='cyan', label='Raw Audio')
|
||||
|
||||
print("offset_index:", offset_index)
|
||||
plt.show()
|
||||
|
||||
|
||||
def export_prov2(
|
||||
gpt_path,
|
||||
vits_path,
|
||||
version,
|
||||
ref_audio_path,
|
||||
ref_text,
|
||||
output_path,
|
||||
device="cpu",
|
||||
is_half=True,
|
||||
lang="auto",
|
||||
):
|
||||
if export_torch_script.sv_cn_model == None:
|
||||
init_sv_cn(device,is_half)
|
||||
|
||||
ref_audio = torch.tensor([load_audio(ref_audio_path, 16000)]).float()
|
||||
ssl = SSLModel()
|
||||
|
||||
print(f"device: {device}")
|
||||
|
||||
ref_seq_id, ref_bert_T, ref_norm_text = get_phones_and_bert(
|
||||
ref_text, lang, "v2"
|
||||
)
|
||||
ref_seq = torch.LongTensor([ref_seq_id]).to(device)
|
||||
ref_bert = ref_bert_T.T
|
||||
if is_half:
|
||||
ref_bert = ref_bert.half()
|
||||
ref_bert = ref_bert.to(ref_seq.device)
|
||||
|
||||
text_seq_id, text_bert_T, norm_text = get_phones_and_bert(
|
||||
"这是一个简单的示例,真没想到这么简单就完成了.The King and His Stories.Once there was a king.He likes to write stories, but his stories were not good.", "auto", "v2"
|
||||
)
|
||||
text_seq = torch.LongTensor([text_seq_id]).to(device)
|
||||
text_bert = text_bert_T.T
|
||||
if is_half:
|
||||
text_bert = text_bert.half()
|
||||
text_bert = text_bert.to(text_seq.device)
|
||||
|
||||
ssl_content = ssl(ref_audio)
|
||||
if is_half:
|
||||
ssl_content = ssl_content.half()
|
||||
ssl_content = ssl_content.to(device)
|
||||
|
||||
sv_model = ExportERes2NetV2(export_torch_script.sv_cn_model)
|
||||
|
||||
# vits_path = "SoVITS_weights_v2/xw_e8_s216.pth"
|
||||
vits = VitsModel(vits_path, version,is_half=is_half,device=device)
|
||||
vits.eval()
|
||||
vits = StepVitsModel(vits, sv_model)
|
||||
|
||||
# gpt_path = "GPT_weights_v2/xw-e15.ckpt"
|
||||
# dict_s1 = torch.load(gpt_path, map_location=device)
|
||||
dict_s1 = torch.load(gpt_path, weights_only=False)
|
||||
raw_t2s = get_raw_t2s_model(dict_s1).to(device)
|
||||
print("#### get_raw_t2s_model ####")
|
||||
print(raw_t2s.config)
|
||||
if is_half:
|
||||
raw_t2s = raw_t2s.half()
|
||||
t2s_m = T2SModel(raw_t2s)
|
||||
t2s_m.eval()
|
||||
# t2s = torch.jit.script(t2s_m).to(device)
|
||||
t2s = t2s_m
|
||||
print("#### script t2s_m ####")
|
||||
|
||||
print("vits.hps.data.sampling_rate:", vits.hps.data.sampling_rate)
|
||||
|
||||
stream_t2s = StreamT2SModel(t2s).to(device)
|
||||
stream_t2s = torch.jit.script(stream_t2s)
|
||||
|
||||
ref_audio_sr = resamplex(ref_audio, 16000, 32000)
|
||||
ref_audio_sr = ref_audio_sr.to(device)
|
||||
if is_half:
|
||||
ref_audio_sr = ref_audio_sr.half()
|
||||
|
||||
top_k = 15
|
||||
|
||||
prompts = vits.extract_latent(ssl_content)
|
||||
|
||||
audio_16k = resamplex(ref_audio_sr, 32000, 16000).to(ref_audio_sr.dtype)
|
||||
sv_emb = sv_model(audio_16k)
|
||||
print("text_seq",text_seq.shape)
|
||||
# torch.jit.trace()
|
||||
|
||||
refer,sv_emb = vits.ref_handle(ref_audio_sr)
|
||||
|
||||
st = time.time()
|
||||
et = time.time()
|
||||
|
||||
y_len, y, xy_pos, k_cache, v_cache = stream_t2s.pre_infer(prompts, ref_seq, text_seq, ref_bert, text_bert, top_k)
|
||||
idx = 1
|
||||
print("y.shape:", y.shape)
|
||||
while True:
|
||||
y, xy_pos, last_token, k_cache, v_cache = stream_t2s(idx, top_k, y_len, y, xy_pos, k_cache, v_cache)
|
||||
# print("y.shape:", y.shape)
|
||||
|
||||
idx+=1
|
||||
# print(idx,'/',1500 , y.shape, y[0,-1].item(), stop)
|
||||
if idx>1500:
|
||||
break
|
||||
|
||||
if last_token == t2s.EOS:
|
||||
break
|
||||
|
||||
at = time.time()
|
||||
print("EOS:",t2s.EOS)
|
||||
|
||||
print(f"frist token: {et - st:.4f} seconds")
|
||||
print(f"all token: {at - st:.4f} seconds")
|
||||
print("sv_emb", sv_emb.shape)
|
||||
print("refer",refer.shape)
|
||||
y = y[:,-idx:].unsqueeze(0)
|
||||
print("y", y.shape)
|
||||
audio = vits(y, text_seq, refer, sv_emb)
|
||||
soundfile.write(f"{output_path}/out_final.wav", audio.float().detach().cpu().numpy(), 32000)
|
||||
|
||||
torch._dynamo.mark_dynamic(ssl_content, 2)
|
||||
torch._dynamo.mark_dynamic(ref_audio_sr, 1)
|
||||
torch._dynamo.mark_dynamic(ref_seq, 1)
|
||||
torch._dynamo.mark_dynamic(text_seq, 1)
|
||||
torch._dynamo.mark_dynamic(ref_bert, 0)
|
||||
torch._dynamo.mark_dynamic(text_bert, 0)
|
||||
torch._dynamo.mark_dynamic(refer, 2)
|
||||
torch._dynamo.mark_dynamic(y, 2)
|
||||
|
||||
inputs = {
|
||||
"forward": (y, text_seq, refer, sv_emb),
|
||||
"extract_latent": ssl_content,
|
||||
"ref_handle": ref_audio_sr,
|
||||
}
|
||||
|
||||
stream_t2s.save(f"{output_path}/t2s.pt")
|
||||
torch.jit.trace_module(vits, inputs=inputs, optimize=True).save(f"{output_path}/vits.pt")
|
||||
torch.jit.script(find_best_audio_offset_fast, optimize=True).save(f"{output_path}/find_best_audio_offset_fast.pt")
|
||||
|
||||
import argparse
|
||||
import os
|
||||
|
||||
if __name__ == "__main__":
|
||||
parser = argparse.ArgumentParser(description="GPT-SoVITS Command Line Tool")
|
||||
parser.add_argument("--gpt_model", required=True, help="Path to the GPT model file")
|
||||
parser.add_argument(
|
||||
"--sovits_model", required=True, help="Path to the SoVITS model file"
|
||||
)
|
||||
parser.add_argument(
|
||||
"--ref_audio", required=True, help="Path to the reference audio file"
|
||||
)
|
||||
parser.add_argument(
|
||||
"--ref_text", required=True, help="Path to the reference text file"
|
||||
)
|
||||
parser.add_argument(
|
||||
"--output_path", required=True, help="Path to the output directory"
|
||||
)
|
||||
parser.add_argument("--device", help="Device to use", default="cuda" if torch.cuda.is_available() else "cpu")
|
||||
parser.add_argument("--version", help="version of the model", default="v2Pro")
|
||||
parser.add_argument("--no-half", action="store_true", help = "Do not use half precision for model weights")
|
||||
parser.add_argument("--lang", default="auto", help="Language for text processing (default: auto)")
|
||||
|
||||
args = parser.parse_args()
|
||||
|
||||
if not os.path.exists(args.output_path):
|
||||
os.makedirs(args.output_path)
|
||||
|
||||
is_half = not args.no_half
|
||||
with torch.no_grad():
|
||||
export_prov2(
|
||||
gpt_path=args.gpt_model,
|
||||
vits_path=args.sovits_model,
|
||||
version=args.version,
|
||||
ref_audio_path=args.ref_audio,
|
||||
ref_text=args.ref_text,
|
||||
output_path=args.output_path,
|
||||
device=args.device,
|
||||
is_half=is_half,
|
||||
lang=args.lang,
|
||||
)
|
||||
+19
-11
@@ -1,24 +1,32 @@
|
||||
import sys,os,torch
|
||||
import sys
|
||||
import os
|
||||
import torch
|
||||
|
||||
sys.path.append(f"{os.getcwd()}/GPT_SoVITS/eres2net")
|
||||
sv_path = "GPT_SoVITS/pretrained_models/sv/pretrained_eres2netv2w24s4ep4.ckpt"
|
||||
from ERes2NetV2 import ERes2NetV2
|
||||
import kaldi as Kaldi
|
||||
|
||||
|
||||
class SV:
|
||||
def __init__(self,device,is_half):
|
||||
pretrained_state = torch.load(sv_path, map_location='cpu', weights_only=False)
|
||||
embedding_model = ERes2NetV2(baseWidth=24,scale=4,expansion=4)
|
||||
def __init__(self, device, is_half):
|
||||
pretrained_state = torch.load(sv_path, map_location="cpu", weights_only=False)
|
||||
embedding_model = ERes2NetV2(baseWidth=24, scale=4, expansion=4)
|
||||
embedding_model.load_state_dict(pretrained_state)
|
||||
embedding_model.eval()
|
||||
self.embedding_model=embedding_model
|
||||
self.embedding_model = embedding_model
|
||||
if is_half == False:
|
||||
self.embedding_model=self.embedding_model.to(device)
|
||||
self.embedding_model = self.embedding_model.to(device)
|
||||
else:
|
||||
self.embedding_model=self.embedding_model.half().to(device)
|
||||
self.is_half=is_half
|
||||
self.embedding_model = self.embedding_model.half().to(device)
|
||||
self.is_half = is_half
|
||||
|
||||
def compute_embedding3(self,wav):
|
||||
def compute_embedding3(self, wav):
|
||||
with torch.no_grad():
|
||||
if self.is_half==True:wav=wav.half()
|
||||
feat = torch.stack([Kaldi.fbank(wav0.unsqueeze(0), num_mel_bins=80, sample_frequency=16000, dither=0) for wav0 in wav])
|
||||
if self.is_half == True:
|
||||
wav = wav.half()
|
||||
feat = torch.stack(
|
||||
[Kaldi.fbank(wav0.unsqueeze(0), num_mel_bins=80, sample_frequency=16000, dither=0) for wav0 in wav]
|
||||
)
|
||||
sv_emb = self.embedding_model.forward3(feat)
|
||||
return sv_emb
|
||||
|
||||
@@ -87,70 +87,116 @@ class LangSegmenter():
|
||||
"en": "en",
|
||||
}
|
||||
|
||||
|
||||
def getTexts(text):
|
||||
def getTexts(text,default_lang = ""):
|
||||
lang_splitter = LangSplitter(lang_map=LangSegmenter.DEFAULT_LANG_MAP)
|
||||
lang_splitter.merge_across_digit = False
|
||||
substr = lang_splitter.split_by_lang(text=text)
|
||||
|
||||
lang_list: list[dict] = []
|
||||
|
||||
have_num = False
|
||||
|
||||
for _, item in enumerate(substr):
|
||||
dict_item = {'lang':item.lang,'text':item.text}
|
||||
|
||||
if dict_item['lang'] == 'digit':
|
||||
if default_lang != "":
|
||||
dict_item['lang'] = default_lang
|
||||
else:
|
||||
have_num = True
|
||||
lang_list = merge_lang(lang_list,dict_item)
|
||||
continue
|
||||
|
||||
# 处理短英文被识别为其他语言的问题
|
||||
if full_en(dict_item['text']):
|
||||
dict_item['lang'] = 'en'
|
||||
lang_list = merge_lang(lang_list,dict_item)
|
||||
continue
|
||||
|
||||
# 处理非日语夹日文的问题(不包含CJK)
|
||||
ja_list: list[dict] = []
|
||||
if dict_item['lang'] != 'ja':
|
||||
ja_list = split_jako('ja',dict_item)
|
||||
if default_lang != "":
|
||||
dict_item['lang'] = default_lang
|
||||
lang_list = merge_lang(lang_list,dict_item)
|
||||
continue
|
||||
else:
|
||||
# 处理非日语夹日文的问题(不包含CJK)
|
||||
ja_list: list[dict] = []
|
||||
if dict_item['lang'] != 'ja':
|
||||
ja_list = split_jako('ja',dict_item)
|
||||
|
||||
if not ja_list:
|
||||
ja_list.append(dict_item)
|
||||
if not ja_list:
|
||||
ja_list.append(dict_item)
|
||||
|
||||
# 处理非韩语夹韩语的问题(不包含CJK)
|
||||
ko_list: list[dict] = []
|
||||
temp_list: list[dict] = []
|
||||
for _, ko_item in enumerate(ja_list):
|
||||
if ko_item["lang"] != 'ko':
|
||||
ko_list = split_jako('ko',ko_item)
|
||||
# 处理非韩语夹韩语的问题(不包含CJK)
|
||||
ko_list: list[dict] = []
|
||||
temp_list: list[dict] = []
|
||||
for _, ko_item in enumerate(ja_list):
|
||||
if ko_item["lang"] != 'ko':
|
||||
ko_list = split_jako('ko',ko_item)
|
||||
|
||||
if ko_list:
|
||||
temp_list.extend(ko_list)
|
||||
else:
|
||||
temp_list.append(ko_item)
|
||||
if ko_list:
|
||||
temp_list.extend(ko_list)
|
||||
else:
|
||||
temp_list.append(ko_item)
|
||||
|
||||
# 未存在非日韩文夹日韩文
|
||||
if len(temp_list) == 1:
|
||||
# 未知语言检查是否为CJK
|
||||
if dict_item['lang'] == 'x':
|
||||
cjk_text = full_cjk(dict_item['text'])
|
||||
if cjk_text:
|
||||
dict_item = {'lang':'zh','text':cjk_text}
|
||||
lang_list = merge_lang(lang_list,dict_item)
|
||||
# 未存在非日韩文夹日韩文
|
||||
if len(temp_list) == 1:
|
||||
# 未知语言检查是否为CJK
|
||||
if dict_item['lang'] == 'x':
|
||||
cjk_text = full_cjk(dict_item['text'])
|
||||
if cjk_text:
|
||||
dict_item = {'lang':'zh','text':cjk_text}
|
||||
lang_list = merge_lang(lang_list,dict_item)
|
||||
else:
|
||||
lang_list = merge_lang(lang_list,dict_item)
|
||||
continue
|
||||
else:
|
||||
lang_list = merge_lang(lang_list,dict_item)
|
||||
continue
|
||||
else:
|
||||
lang_list = merge_lang(lang_list,dict_item)
|
||||
continue
|
||||
continue
|
||||
|
||||
# 存在非日韩文夹日韩文
|
||||
for _, temp_item in enumerate(temp_list):
|
||||
# 未知语言检查是否为CJK
|
||||
if temp_item['lang'] == 'x':
|
||||
cjk_text = full_cjk(dict_item['text'])
|
||||
if cjk_text:
|
||||
dict_item = {'lang':'zh','text':cjk_text}
|
||||
lang_list = merge_lang(lang_list,dict_item)
|
||||
# 存在非日韩文夹日韩文
|
||||
for _, temp_item in enumerate(temp_list):
|
||||
# 未知语言检查是否为CJK
|
||||
if temp_item['lang'] == 'x':
|
||||
cjk_text = full_cjk(temp_item['text'])
|
||||
if cjk_text:
|
||||
lang_list = merge_lang(lang_list,{'lang':'zh','text':cjk_text})
|
||||
else:
|
||||
lang_list = merge_lang(lang_list,temp_item)
|
||||
else:
|
||||
lang_list = merge_lang(lang_list,dict_item)
|
||||
else:
|
||||
lang_list = merge_lang(lang_list,temp_item)
|
||||
lang_list = merge_lang(lang_list,temp_item)
|
||||
|
||||
# 有数字
|
||||
if have_num:
|
||||
temp_list = lang_list
|
||||
lang_list = []
|
||||
for i, temp_item in enumerate(temp_list):
|
||||
if temp_item['lang'] == 'digit':
|
||||
if default_lang:
|
||||
temp_item['lang'] = default_lang
|
||||
elif lang_list and i == len(temp_list) - 1:
|
||||
temp_item['lang'] = lang_list[-1]['lang']
|
||||
elif not lang_list and i < len(temp_list) - 1:
|
||||
temp_item['lang'] = temp_list[1]['lang']
|
||||
elif lang_list and i < len(temp_list) - 1:
|
||||
if lang_list[-1]['lang'] == temp_list[i + 1]['lang']:
|
||||
temp_item['lang'] = lang_list[-1]['lang']
|
||||
elif lang_list[-1]['text'][-1] in [",",".","!","?",",","。","!","?"]:
|
||||
temp_item['lang'] = temp_list[i + 1]['lang']
|
||||
elif temp_list[i + 1]['text'][0] in [",",".","!","?",",","。","!","?"]:
|
||||
temp_item['lang'] = lang_list[-1]['lang']
|
||||
elif temp_item['text'][-1] in ["。","."]:
|
||||
temp_item['lang'] = lang_list[-1]['lang']
|
||||
elif len(lang_list[-1]['text']) >= len(temp_list[i + 1]['text']):
|
||||
temp_item['lang'] = lang_list[-1]['lang']
|
||||
else:
|
||||
temp_item['lang'] = temp_list[i + 1]['lang']
|
||||
else:
|
||||
temp_item['lang'] = 'zh'
|
||||
|
||||
lang_list = merge_lang(lang_list,temp_item)
|
||||
|
||||
|
||||
# 筛X
|
||||
temp_list = lang_list
|
||||
lang_list = []
|
||||
for _, temp_item in enumerate(temp_list):
|
||||
@@ -172,4 +218,8 @@ if __name__ == "__main__":
|
||||
print(LangSegmenter.getTexts(text))
|
||||
|
||||
text = "ねえ、知ってる?最近、僕は天文学を勉強してるんだ。君の瞳が星空みたいにキラキラしてるからさ。"
|
||||
print(LangSegmenter.getTexts(text))
|
||||
|
||||
text = "当时ThinkPad T60刚刚发布,一同推出的还有一款名为Advanced Dock的扩展坞配件。这款扩展坞通过连接T60底部的插槽,扩展出包括PCIe在内的一大堆接口,并且自带电源,让T60可以安装桌面显卡来提升性能。"
|
||||
print(LangSegmenter.getTexts(text,"zh"))
|
||||
print(LangSegmenter.getTexts(text))
|
||||
@@ -181,20 +181,6 @@ def text_normalize(text):
|
||||
return dest_text
|
||||
|
||||
|
||||
# 不排除英文的文本格式化
|
||||
def mix_text_normalize(text):
|
||||
# https://github.com/PaddlePaddle/PaddleSpeech/tree/develop/paddlespeech/t2s/frontend/zh_normalization
|
||||
tx = TextNormalizer()
|
||||
sentences = tx.normalize(text)
|
||||
dest_text = ""
|
||||
for sentence in sentences:
|
||||
dest_text += replace_punctuation_with_en(sentence)
|
||||
|
||||
# 避免重复标点引起的参考泄露
|
||||
dest_text = replace_consecutive_punctuation(dest_text)
|
||||
return dest_text
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
text = "啊——但是《原神》是由,米哈\游自主,研发的一款全.新开放世界.冒险游戏"
|
||||
text = "呣呣呣~就是…大人的鼹鼠党吧?"
|
||||
|
||||
+12
-19
@@ -180,10 +180,15 @@ def _merge_erhua(initials: list[str], finals: list[str], word: str, pos: str) ->
|
||||
def _g2p(segments):
|
||||
phones_list = []
|
||||
word2ph = []
|
||||
for seg in segments:
|
||||
g2pw_batch_results = []
|
||||
g2pw_batch_cursor = 0
|
||||
processed_segments = [re.sub("[a-zA-Z]+", "", seg) for seg in segments]
|
||||
if is_g2pw:
|
||||
batch_inputs = [seg for seg in processed_segments if seg]
|
||||
g2pw_batch_results = g2pw._g2pw(batch_inputs) if batch_inputs else []
|
||||
|
||||
for seg in processed_segments:
|
||||
pinyins = []
|
||||
# Replace all English words in the sentence
|
||||
seg = re.sub("[a-zA-Z]+", "", seg)
|
||||
seg_cut = psg.lcut(seg)
|
||||
seg_cut = tone_modifier.pre_merge_for_modify(seg_cut)
|
||||
initials = []
|
||||
@@ -204,8 +209,10 @@ def _g2p(segments):
|
||||
finals = sum(finals, [])
|
||||
print("pypinyin结果", initials, finals)
|
||||
else:
|
||||
# g2pw采用整句推理
|
||||
pinyins = g2pw.lazy_pinyin(seg, neutral_tone_with_five=True, style=Style.TONE3)
|
||||
# g2pw采用整句推理(批量推理,逐句取结果)
|
||||
if seg:
|
||||
pinyins = g2pw_batch_results[g2pw_batch_cursor]
|
||||
g2pw_batch_cursor += 1
|
||||
|
||||
pre_word_length = 0
|
||||
for word, pos in seg_cut:
|
||||
@@ -326,20 +333,6 @@ def text_normalize(text):
|
||||
return dest_text
|
||||
|
||||
|
||||
# 不排除英文的文本格式化
|
||||
def mix_text_normalize(text):
|
||||
# https://github.com/PaddlePaddle/PaddleSpeech/tree/develop/paddlespeech/t2s/frontend/zh_normalization
|
||||
tx = TextNormalizer()
|
||||
sentences = tx.normalize(text)
|
||||
dest_text = ""
|
||||
for sentence in sentences:
|
||||
dest_text += replace_punctuation_with_en(sentence)
|
||||
|
||||
# 避免重复标点引起的参考泄露
|
||||
dest_text = replace_consecutive_punctuation(dest_text)
|
||||
return dest_text
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
text = "啊——但是《原神》是由,米哈\游自主,研发的一款全.新开放世界.冒险游戏"
|
||||
text = "呣呣呣~就是…大人的鼹鼠党吧?"
|
||||
|
||||
@@ -238,6 +238,46 @@ def _expand_number(m):
|
||||
return _inflect.number_to_words(num, andword="")
|
||||
|
||||
|
||||
# 加减乘除
|
||||
RE_ASMD = re.compile(
|
||||
r"((-?)((\d+)(\.\d+)?[⁰¹²³⁴⁵⁶⁷⁸⁹ˣʸⁿ]*)|(\.\d+[⁰¹²³⁴⁵⁶⁷⁸⁹ˣʸⁿ]*)|([A-Za-z][⁰¹²³⁴⁵⁶⁷⁸⁹ˣʸⁿ]*))\s+([\+\-\×÷=])\s+((-?)((\d+)(\.\d+)?[⁰¹²³⁴⁵⁶⁷⁸⁹ˣʸⁿ]*)|(\.\d+[⁰¹²³⁴⁵⁶⁷⁸⁹ˣʸⁿ]*)|([A-Za-z][⁰¹²³⁴⁵⁶⁷⁸⁹ˣʸⁿ]*))"
|
||||
)
|
||||
# RE_ASMD = re.compile(
|
||||
# r"\b((-?)((\d+)(\.\d+)?[⁰¹²³⁴⁵⁶⁷⁸⁹ˣʸⁿ]*)|(\.\d+[⁰¹²³⁴⁵⁶⁷⁸⁹ˣʸⁿ]*)|([A-Za-z][⁰¹²³⁴⁵⁶⁷⁸⁹ˣʸⁿ]*))([\+\-\×÷=])((-?)((\d+)(\.\d+)?[⁰¹²³⁴⁵⁶⁷⁸⁹ˣʸⁿ]*)|(\.\d+[⁰¹²³⁴⁵⁶⁷⁸⁹ˣʸⁿ]*)|([A-Za-z][⁰¹²³⁴⁵⁶⁷⁸⁹ˣʸⁿ]*))\b"
|
||||
# )
|
||||
|
||||
asmd_map = {"+": " plus ", "-": " minus ", "×": " times ", "÷": " divided by ", "=": " Equals "}
|
||||
|
||||
|
||||
def replace_asmd(match) -> str:
|
||||
"""
|
||||
Args:
|
||||
match (re.Match)
|
||||
Returns:
|
||||
str
|
||||
"""
|
||||
result = match.group(1) + asmd_map[match.group(8)] + match.group(9)
|
||||
return result
|
||||
|
||||
|
||||
RE_INTEGER = re.compile(r"(?:^|\s+)(-)" r"(\d+)")
|
||||
|
||||
|
||||
def replace_negative_num(match) -> str:
|
||||
"""
|
||||
Args:
|
||||
match (re.Match)
|
||||
Returns:
|
||||
str
|
||||
"""
|
||||
sign = match.group(1)
|
||||
number = match.group(2)
|
||||
sign: str = "negative " if sign else ""
|
||||
result = f"{sign}{number}"
|
||||
return result
|
||||
|
||||
|
||||
|
||||
def normalize(text):
|
||||
"""
|
||||
!!! 所有的处理都需要正确的输入 !!!
|
||||
@@ -245,7 +285,13 @@ def normalize(text):
|
||||
"""
|
||||
|
||||
text = re.sub(_ordinal_number_re, _convert_ordinal, text)
|
||||
text = re.sub(r"(?<!\d)-|-(?!\d)", " minus ", text)
|
||||
|
||||
# 处理数学运算
|
||||
# 替换text = re.sub(r"(?<!\d)-|-(?!\d)", " minus ", text)
|
||||
while RE_ASMD.search(text):
|
||||
text = RE_ASMD.sub(replace_asmd, text)
|
||||
text = RE_INTEGER.sub(replace_negative_num, text)
|
||||
|
||||
text = re.sub(_comma_number_re, _remove_commas, text)
|
||||
text = re.sub(_time_re, _expand_time, text)
|
||||
text = re.sub(_measurement_re, _expand_measurement, text)
|
||||
|
||||
@@ -18,6 +18,7 @@ Credits
|
||||
|
||||
from typing import Dict
|
||||
from typing import List
|
||||
from typing import Optional
|
||||
from typing import Tuple
|
||||
|
||||
import numpy as np
|
||||
@@ -37,6 +38,8 @@ def prepare_onnx_input(
|
||||
use_mask: bool = False,
|
||||
window_size: int = None,
|
||||
max_len: int = 512,
|
||||
char2id: Optional[Dict[str, int]] = None,
|
||||
char_phoneme_masks: Optional[Dict[str, List[int]]] = None,
|
||||
) -> Dict[str, np.array]:
|
||||
if window_size is not None:
|
||||
truncated_texts, truncated_query_ids = _truncate_texts(
|
||||
@@ -48,33 +51,88 @@ def prepare_onnx_input(
|
||||
phoneme_masks = []
|
||||
char_ids = []
|
||||
position_ids = []
|
||||
tokenized_cache = {}
|
||||
|
||||
if char2id is None:
|
||||
char2id = {char: idx for idx, char in enumerate(chars)}
|
||||
if use_mask:
|
||||
if char_phoneme_masks is None:
|
||||
char_phoneme_masks = {
|
||||
char: [1 if i in char2phonemes[char] else 0 for i in range(len(labels))]
|
||||
for char in char2phonemes
|
||||
}
|
||||
else:
|
||||
full_phoneme_mask = [1] * len(labels)
|
||||
|
||||
for idx in range(len(texts)):
|
||||
text = (truncated_texts if window_size else texts)[idx].lower()
|
||||
query_id = (truncated_query_ids if window_size else query_ids)[idx]
|
||||
|
||||
try:
|
||||
tokens, text2token, token2text = tokenize_and_map(tokenizer=tokenizer, text=text)
|
||||
except Exception:
|
||||
print(f'warning: text "{text}" is invalid')
|
||||
return {}
|
||||
cached = tokenized_cache.get(text)
|
||||
if cached is None:
|
||||
try:
|
||||
tokens, text2token, token2text = tokenize_and_map(tokenizer=tokenizer, text=text)
|
||||
except Exception:
|
||||
print(f'warning: text "{text}" is invalid')
|
||||
return {}
|
||||
|
||||
text, query_id, tokens, text2token, token2text = _truncate(
|
||||
max_len=max_len, text=text, query_id=query_id, tokens=tokens, text2token=text2token, token2text=token2text
|
||||
)
|
||||
if len(tokens) <= max_len - 2:
|
||||
processed_tokens = ["[CLS]"] + tokens + ["[SEP]"]
|
||||
shared_input_id = list(np.array(tokenizer.convert_tokens_to_ids(processed_tokens)))
|
||||
shared_token_type_id = list(np.zeros((len(processed_tokens),), dtype=int))
|
||||
shared_attention_mask = list(np.ones((len(processed_tokens),), dtype=int))
|
||||
cached = {
|
||||
"is_short": True,
|
||||
"tokens": tokens,
|
||||
"text2token": text2token,
|
||||
"token2text": token2text,
|
||||
"input_id": shared_input_id,
|
||||
"token_type_id": shared_token_type_id,
|
||||
"attention_mask": shared_attention_mask,
|
||||
}
|
||||
else:
|
||||
cached = {
|
||||
"is_short": False,
|
||||
"tokens": tokens,
|
||||
"text2token": text2token,
|
||||
"token2text": token2text,
|
||||
}
|
||||
tokenized_cache[text] = cached
|
||||
|
||||
processed_tokens = ["[CLS]"] + tokens + ["[SEP]"]
|
||||
if cached["is_short"]:
|
||||
text_for_query = text
|
||||
query_id_for_query = query_id
|
||||
text2token_for_query = cached["text2token"]
|
||||
input_id = cached["input_id"]
|
||||
token_type_id = cached["token_type_id"]
|
||||
attention_mask = cached["attention_mask"]
|
||||
else:
|
||||
(
|
||||
text_for_query,
|
||||
query_id_for_query,
|
||||
tokens_for_query,
|
||||
text2token_for_query,
|
||||
_token2text_for_query,
|
||||
) = _truncate(
|
||||
max_len=max_len,
|
||||
text=text,
|
||||
query_id=query_id,
|
||||
tokens=cached["tokens"],
|
||||
text2token=cached["text2token"],
|
||||
token2text=cached["token2text"],
|
||||
)
|
||||
processed_tokens = ["[CLS]"] + tokens_for_query + ["[SEP]"]
|
||||
input_id = list(np.array(tokenizer.convert_tokens_to_ids(processed_tokens)))
|
||||
token_type_id = list(np.zeros((len(processed_tokens),), dtype=int))
|
||||
attention_mask = list(np.ones((len(processed_tokens),), dtype=int))
|
||||
|
||||
input_id = list(np.array(tokenizer.convert_tokens_to_ids(processed_tokens)))
|
||||
token_type_id = list(np.zeros((len(processed_tokens),), dtype=int))
|
||||
attention_mask = list(np.ones((len(processed_tokens),), dtype=int))
|
||||
|
||||
query_char = text[query_id]
|
||||
phoneme_mask = (
|
||||
[1 if i in char2phonemes[query_char] else 0 for i in range(len(labels))] if use_mask else [1] * len(labels)
|
||||
)
|
||||
char_id = chars.index(query_char)
|
||||
position_id = text2token[query_id] + 1 # [CLS] token locate at first place
|
||||
query_char = text_for_query[query_id_for_query]
|
||||
if use_mask:
|
||||
phoneme_mask = char_phoneme_masks[query_char]
|
||||
else:
|
||||
phoneme_mask = full_phoneme_mask
|
||||
char_id = char2id[query_char]
|
||||
position_id = text2token_for_query[query_id_for_query] + 1 # [CLS] token locate at first place
|
||||
|
||||
input_ids.append(input_id)
|
||||
token_type_ids.append(token_type_id)
|
||||
@@ -83,10 +141,15 @@ def prepare_onnx_input(
|
||||
char_ids.append(char_id)
|
||||
position_ids.append(position_id)
|
||||
|
||||
max_token_length = max(len(seq) for seq in input_ids)
|
||||
|
||||
def _pad_sequences(sequences, pad_value=0):
|
||||
return [seq + [pad_value] * (max_token_length - len(seq)) for seq in sequences]
|
||||
|
||||
outputs = {
|
||||
"input_ids": np.array(input_ids).astype(np.int64),
|
||||
"token_type_ids": np.array(token_type_ids).astype(np.int64),
|
||||
"attention_masks": np.array(attention_masks).astype(np.int64),
|
||||
"input_ids": np.array(_pad_sequences(input_ids, pad_value=0)).astype(np.int64),
|
||||
"token_type_ids": np.array(_pad_sequences(token_type_ids, pad_value=0)).astype(np.int64),
|
||||
"attention_masks": np.array(_pad_sequences(attention_masks, pad_value=0)).astype(np.int64),
|
||||
"phoneme_masks": np.array(phoneme_masks).astype(np.float32),
|
||||
"char_ids": np.array(char_ids).astype(np.int64),
|
||||
"position_ids": np.array(position_ids).astype(np.int64),
|
||||
|
||||
@@ -1,5 +1,6 @@
|
||||
# This code is modified from https://github.com/mozillazg/pypinyin-g2pW
|
||||
|
||||
import hashlib
|
||||
import pickle
|
||||
import os
|
||||
|
||||
@@ -14,6 +15,16 @@ current_file_path = os.path.dirname(__file__)
|
||||
CACHE_PATH = os.path.join(current_file_path, "polyphonic.pickle")
|
||||
PP_DICT_PATH = os.path.join(current_file_path, "polyphonic.rep")
|
||||
PP_FIX_DICT_PATH = os.path.join(current_file_path, "polyphonic-fix.rep")
|
||||
MD5_PATH = os.path.join(current_file_path, "polyphonic.md5")
|
||||
|
||||
def get_file_md5(file_path):
|
||||
if not os.path.exists(file_path):
|
||||
return ""
|
||||
hasher = hashlib.md5()
|
||||
with open(file_path, "rb") as f:
|
||||
for chunk in iter(lambda: f.read(4096), b""):
|
||||
hasher.update(chunk)
|
||||
return hasher.hexdigest()
|
||||
|
||||
|
||||
class G2PWPinyin(Pinyin):
|
||||
@@ -115,13 +126,22 @@ def cache_dict(polyphonic_dict, file_path):
|
||||
|
||||
|
||||
def get_dict():
|
||||
if os.path.exists(CACHE_PATH):
|
||||
new_md5 = get_file_md5(PP_DICT_PATH) + get_file_md5(PP_FIX_DICT_PATH)
|
||||
old_md5 = ""
|
||||
if os.path.exists(MD5_PATH):
|
||||
with open(MD5_PATH, "r", encoding="utf-8") as f:
|
||||
old_md5 = f.read().strip()
|
||||
need_rebuild = (not os.path.exists(CACHE_PATH)) or (new_md5 != old_md5)
|
||||
|
||||
if not need_rebuild:
|
||||
with open(CACHE_PATH, "rb") as pickle_file:
|
||||
polyphonic_dict = pickle.load(pickle_file)
|
||||
else:
|
||||
print("Rebuilding Polyphonic Dictionary: " + f"{old_md5} -> {new_md5}")
|
||||
polyphonic_dict = read_dict()
|
||||
cache_dict(polyphonic_dict, CACHE_PATH)
|
||||
|
||||
with open(MD5_PATH, "w", encoding="utf-8") as f:
|
||||
f.write(new_md5)
|
||||
return polyphonic_dict
|
||||
|
||||
|
||||
|
||||
@@ -3,7 +3,6 @@
|
||||
|
||||
import json
|
||||
import os
|
||||
import traceback
|
||||
import warnings
|
||||
import zipfile
|
||||
from typing import Any, Dict, List, Tuple
|
||||
@@ -11,7 +10,6 @@ from typing import Any, Dict, List, Tuple
|
||||
import numpy as np
|
||||
import onnxruntime
|
||||
import requests
|
||||
import torch
|
||||
from opencc import OpenCC
|
||||
from pypinyin import Style, pinyin
|
||||
from transformers.models.auto.tokenization_auto import AutoTokenizer
|
||||
@@ -23,8 +21,8 @@ from .utils import load_config
|
||||
onnxruntime.set_default_logger_severity(3)
|
||||
try:
|
||||
onnxruntime.preload_dlls()
|
||||
except:pass
|
||||
#traceback.print_exc()
|
||||
except Exception:
|
||||
pass
|
||||
warnings.filterwarnings("ignore")
|
||||
|
||||
model_version = "1.1"
|
||||
@@ -55,6 +53,24 @@ def predict(session, onnx_input: Dict[str, Any], labels: List[str]) -> Tuple[Lis
|
||||
return all_preds, all_confidences
|
||||
|
||||
|
||||
def _load_json_from_candidates(filename: str, candidate_dirs: List[str]) -> Dict[str, Any]:
|
||||
for candidate_dir in candidate_dirs:
|
||||
if not candidate_dir:
|
||||
continue
|
||||
json_path = os.path.join(candidate_dir, filename)
|
||||
if os.path.exists(json_path):
|
||||
with open(json_path, "r", encoding="utf-8") as fr:
|
||||
return json.load(fr)
|
||||
raise FileNotFoundError(f"Cannot locate {filename} in candidate dirs: {candidate_dirs}")
|
||||
|
||||
|
||||
def _find_first_existing_file(*paths: str) -> str:
|
||||
for path in paths:
|
||||
if path and os.path.exists(path):
|
||||
return path
|
||||
raise FileNotFoundError(f"Files not found: {paths}")
|
||||
|
||||
|
||||
def download_and_decompress(model_dir: str = "G2PWModel/"):
|
||||
if not os.path.exists(model_dir):
|
||||
parent_directory = os.path.dirname(model_dir)
|
||||
@@ -62,7 +78,7 @@ def download_and_decompress(model_dir: str = "G2PWModel/"):
|
||||
extract_dir = os.path.join(parent_directory, "G2PWModel_1.1")
|
||||
extract_dir_new = os.path.join(parent_directory, "G2PWModel")
|
||||
print("Downloading g2pw model...")
|
||||
modelscope_url = "https://www.modelscope.cn/models/kamiorinn/g2pw/resolve/master/G2PWModel_1.1.zip" # "https://paddlespeech.cdn.bcebos.com/Parakeet/released_models/g2p/G2PWModel_1.1.zip"
|
||||
modelscope_url = "https://www.modelscope.cn/models/kamiorinn/g2pw/resolve/master/G2PWModel_1.1.zip"
|
||||
with requests.get(modelscope_url, stream=True) as r:
|
||||
r.raise_for_status()
|
||||
with open(zip_dir, "wb") as f:
|
||||
@@ -79,7 +95,7 @@ def download_and_decompress(model_dir: str = "G2PWModel/"):
|
||||
return model_dir
|
||||
|
||||
|
||||
class G2PWOnnxConverter:
|
||||
class _G2PWBaseOnnxConverter:
|
||||
def __init__(
|
||||
self,
|
||||
model_dir: str = "G2PWModel/",
|
||||
@@ -87,33 +103,16 @@ class G2PWOnnxConverter:
|
||||
model_source: str = None,
|
||||
enable_non_tradional_chinese: bool = False,
|
||||
):
|
||||
uncompress_path = download_and_decompress(model_dir)
|
||||
|
||||
sess_options = onnxruntime.SessionOptions()
|
||||
sess_options.graph_optimization_level = onnxruntime.GraphOptimizationLevel.ORT_ENABLE_ALL
|
||||
sess_options.execution_mode = onnxruntime.ExecutionMode.ORT_SEQUENTIAL
|
||||
sess_options.intra_op_num_threads = 2 if torch.cuda.is_available() else 0
|
||||
try:
|
||||
self.session_g2pW = onnxruntime.InferenceSession(
|
||||
os.path.join(uncompress_path, "g2pW.onnx"),
|
||||
sess_options=sess_options,
|
||||
providers=["CUDAExecutionProvider", "CPUExecutionProvider"],
|
||||
)
|
||||
except:
|
||||
self.session_g2pW = onnxruntime.InferenceSession(
|
||||
os.path.join(uncompress_path, "g2pW.onnx"),
|
||||
sess_options=sess_options,
|
||||
providers=["CPUExecutionProvider"],
|
||||
)
|
||||
self.config = load_config(config_path=os.path.join(uncompress_path, "config.py"), use_default=True)
|
||||
self.model_dir = download_and_decompress(model_dir)
|
||||
self.config = load_config(config_path=os.path.join(self.model_dir, "config.py"), use_default=True)
|
||||
|
||||
self.model_source = model_source if model_source else self.config.model_source
|
||||
self.enable_opencc = enable_non_tradional_chinese
|
||||
|
||||
self.tokenizer = AutoTokenizer.from_pretrained(self.model_source)
|
||||
|
||||
polyphonic_chars_path = os.path.join(uncompress_path, "POLYPHONIC_CHARS.txt")
|
||||
monophonic_chars_path = os.path.join(uncompress_path, "MONOPHONIC_CHARS.txt")
|
||||
polyphonic_chars_path = os.path.join(self.model_dir, "POLYPHONIC_CHARS.txt")
|
||||
monophonic_chars_path = os.path.join(self.model_dir, "MONOPHONIC_CHARS.txt")
|
||||
|
||||
self.polyphonic_chars = [
|
||||
line.split("\t") for line in open(polyphonic_chars_path, encoding="utf-8").read().strip().split("\n")
|
||||
]
|
||||
@@ -149,31 +148,47 @@ class G2PWOnnxConverter:
|
||||
)
|
||||
|
||||
self.chars = sorted(list(self.char2phonemes.keys()))
|
||||
self.char2id = {char: idx for idx, char in enumerate(self.chars)}
|
||||
self.char_phoneme_masks = (
|
||||
{
|
||||
char: [1 if i in self.char2phonemes[char] else 0 for i in range(len(self.labels))]
|
||||
for char in self.char2phonemes
|
||||
}
|
||||
if self.config.use_mask
|
||||
else None
|
||||
)
|
||||
|
||||
self.polyphonic_chars_new = set(self.chars)
|
||||
for char in self.non_polyphonic:
|
||||
if char in self.polyphonic_chars_new:
|
||||
self.polyphonic_chars_new.remove(char)
|
||||
self.polyphonic_chars_new.discard(char)
|
||||
|
||||
self.monophonic_chars_dict = {char: phoneme for char, phoneme in self.monophonic_chars}
|
||||
for char in self.non_monophonic:
|
||||
if char in self.monophonic_chars_dict:
|
||||
self.monophonic_chars_dict.pop(char)
|
||||
self.monophonic_chars_dict.pop(char, None)
|
||||
|
||||
self.pos_tags = ["UNK", "A", "C", "D", "I", "N", "P", "T", "V", "DE", "SHI"]
|
||||
default_asset_dir = os.path.normpath(os.path.join(os.path.dirname(__file__), "..", "G2PWModel"))
|
||||
candidate_asset_dirs = [self.model_dir, default_asset_dir]
|
||||
self.bopomofo_convert_dict = _load_json_from_candidates(
|
||||
"bopomofo_to_pinyin_wo_tune_dict.json", candidate_asset_dirs
|
||||
)
|
||||
self.char_bopomofo_dict = _load_json_from_candidates("char_bopomofo_dict.json", candidate_asset_dirs)
|
||||
|
||||
with open(os.path.join(uncompress_path, "bopomofo_to_pinyin_wo_tune_dict.json"), "r", encoding="utf-8") as fr:
|
||||
self.bopomofo_convert_dict = json.load(fr)
|
||||
self.style_convert_func = {
|
||||
"bopomofo": lambda x: x,
|
||||
"pinyin": self._convert_bopomofo_to_pinyin,
|
||||
}[style]
|
||||
|
||||
with open(os.path.join(uncompress_path, "char_bopomofo_dict.json"), "r", encoding="utf-8") as fr:
|
||||
self.char_bopomofo_dict = json.load(fr)
|
||||
|
||||
if self.enable_opencc:
|
||||
self.cc = OpenCC("s2tw")
|
||||
self.enable_sentence_dedup = os.getenv("g2pw_sentence_dedup", "true").strip().lower() in {
|
||||
"1",
|
||||
"true",
|
||||
"yes",
|
||||
"y",
|
||||
"on",
|
||||
}
|
||||
# 聚焦到多音字附近上下文,默认左右各16字;设为0表示关闭裁剪(整句)。
|
||||
self.polyphonic_context_chars = max(0, int(os.getenv("g2pw_polyphonic_context_chars", "16")))
|
||||
|
||||
def _convert_bopomofo_to_pinyin(self, bopomofo: str) -> str:
|
||||
tone = bopomofo[-1]
|
||||
@@ -181,9 +196,8 @@ class G2PWOnnxConverter:
|
||||
component = self.bopomofo_convert_dict.get(bopomofo[:-1])
|
||||
if component:
|
||||
return component + tone
|
||||
else:
|
||||
print(f'Warning: "{bopomofo}" cannot convert to pinyin')
|
||||
return None
|
||||
print(f'Warning: "{bopomofo}" cannot convert to pinyin')
|
||||
return None
|
||||
|
||||
def __call__(self, sentences: List[str]) -> List[List[str]]:
|
||||
if isinstance(sentences, str):
|
||||
@@ -197,51 +211,147 @@ class G2PWOnnxConverter:
|
||||
translated_sentences.append(translated_sent)
|
||||
sentences = translated_sentences
|
||||
|
||||
texts, query_ids, sent_ids, partial_results = self._prepare_data(sentences=sentences)
|
||||
texts, model_query_ids, result_query_ids, sent_ids, partial_results = self._prepare_data(sentences=sentences)
|
||||
if len(texts) == 0:
|
||||
# sentences no polyphonic words
|
||||
return partial_results
|
||||
|
||||
onnx_input = prepare_onnx_input(
|
||||
model_input = prepare_onnx_input(
|
||||
tokenizer=self.tokenizer,
|
||||
labels=self.labels,
|
||||
char2phonemes=self.char2phonemes,
|
||||
chars=self.chars,
|
||||
texts=texts,
|
||||
query_ids=query_ids,
|
||||
query_ids=model_query_ids,
|
||||
use_mask=self.config.use_mask,
|
||||
window_size=None,
|
||||
char2id=self.char2id,
|
||||
char_phoneme_masks=self.char_phoneme_masks,
|
||||
)
|
||||
|
||||
preds, confidences = predict(session=self.session_g2pW, onnx_input=onnx_input, labels=self.labels)
|
||||
if not model_input:
|
||||
return partial_results
|
||||
|
||||
if self.enable_sentence_dedup:
|
||||
preds, _confidences = self._predict_with_sentence_dedup(model_input=model_input, texts=texts)
|
||||
else:
|
||||
preds, _confidences = self._predict(model_input=model_input)
|
||||
|
||||
if self.config.use_char_phoneme:
|
||||
preds = [pred.split(" ")[1] for pred in preds]
|
||||
|
||||
results = partial_results
|
||||
for sent_id, query_id, pred in zip(sent_ids, query_ids, preds):
|
||||
for sent_id, query_id, pred in zip(sent_ids, result_query_ids, preds):
|
||||
results[sent_id][query_id] = self.style_convert_func(pred)
|
||||
|
||||
return results
|
||||
|
||||
def _prepare_data(self, sentences: List[str]) -> Tuple[List[str], List[int], List[int], List[List[str]]]:
|
||||
texts, query_ids, sent_ids, partial_results = [], [], [], []
|
||||
def _prepare_data(
|
||||
self, sentences: List[str]
|
||||
) -> Tuple[List[str], List[int], List[int], List[int], List[List[str]]]:
|
||||
texts, model_query_ids, result_query_ids, sent_ids, partial_results = [], [], [], [], []
|
||||
for sent_id, sent in enumerate(sentences):
|
||||
# pypinyin works well for Simplified Chinese than Traditional Chinese
|
||||
sent_s = tranditional_to_simplified(sent)
|
||||
pypinyin_result = pinyin(sent_s, neutral_tone_with_five=True, style=Style.TONE3)
|
||||
partial_result = [None] * len(sent)
|
||||
polyphonic_indices: List[int] = []
|
||||
for i, char in enumerate(sent):
|
||||
if char in self.polyphonic_chars_new:
|
||||
texts.append(sent)
|
||||
query_ids.append(i)
|
||||
sent_ids.append(sent_id)
|
||||
polyphonic_indices.append(i)
|
||||
elif char in self.monophonic_chars_dict:
|
||||
partial_result[i] = self.style_convert_func(self.monophonic_chars_dict[char])
|
||||
elif char in self.char_bopomofo_dict:
|
||||
partial_result[i] = pypinyin_result[i][0]
|
||||
# partial_result[i] = self.style_convert_func(self.char_bopomofo_dict[char][0])
|
||||
else:
|
||||
partial_result[i] = pypinyin_result[i][0]
|
||||
|
||||
if polyphonic_indices:
|
||||
if self.polyphonic_context_chars > 0:
|
||||
left = max(0, polyphonic_indices[0] - self.polyphonic_context_chars)
|
||||
right = min(len(sent), polyphonic_indices[-1] + self.polyphonic_context_chars + 1)
|
||||
sent_for_predict = sent[left:right]
|
||||
query_offset = left
|
||||
else:
|
||||
sent_for_predict = sent
|
||||
query_offset = 0
|
||||
|
||||
for index in polyphonic_indices:
|
||||
texts.append(sent_for_predict)
|
||||
model_query_ids.append(index - query_offset)
|
||||
result_query_ids.append(index)
|
||||
sent_ids.append(sent_id)
|
||||
|
||||
partial_results.append(partial_result)
|
||||
return texts, query_ids, sent_ids, partial_results
|
||||
return texts, model_query_ids, result_query_ids, sent_ids, partial_results
|
||||
|
||||
def _predict(self, model_input: Dict[str, Any]) -> Tuple[List[str], List[float]]:
|
||||
raise NotImplementedError
|
||||
|
||||
def _predict_with_sentence_dedup(
|
||||
self, model_input: Dict[str, Any], texts: List[str]
|
||||
) -> Tuple[List[str], List[float]]:
|
||||
if len(texts) <= 1:
|
||||
return self._predict(model_input=model_input)
|
||||
|
||||
grouped_indices: Dict[str, List[int]] = {}
|
||||
for idx, text in enumerate(texts):
|
||||
grouped_indices.setdefault(text, []).append(idx)
|
||||
|
||||
if all(len(indices) == 1 for indices in grouped_indices.values()):
|
||||
return self._predict(model_input=model_input)
|
||||
|
||||
preds: List[str] = [""] * len(texts)
|
||||
confidences: List[float] = [0.0] * len(texts)
|
||||
for indices in grouped_indices.values():
|
||||
group_input = {name: value[indices] for name, value in model_input.items()}
|
||||
if len(indices) > 1:
|
||||
for name in ("input_ids", "token_type_ids", "attention_masks"):
|
||||
group_input[name] = group_input[name][:1]
|
||||
|
||||
group_preds, group_confidences = self._predict(model_input=group_input)
|
||||
for output_idx, pred, confidence in zip(indices, group_preds, group_confidences):
|
||||
preds[output_idx] = pred
|
||||
confidences[output_idx] = confidence
|
||||
|
||||
return preds, confidences
|
||||
|
||||
|
||||
class G2PWOnnxConverter(_G2PWBaseOnnxConverter):
|
||||
def __init__(
|
||||
self,
|
||||
model_dir: str = "G2PWModel/",
|
||||
style: str = "bopomofo",
|
||||
model_source: str = None,
|
||||
enable_non_tradional_chinese: bool = False,
|
||||
):
|
||||
super().__init__(
|
||||
model_dir=model_dir,
|
||||
style=style,
|
||||
model_source=model_source,
|
||||
enable_non_tradional_chinese=enable_non_tradional_chinese,
|
||||
)
|
||||
|
||||
sess_options = onnxruntime.SessionOptions()
|
||||
sess_options.graph_optimization_level = onnxruntime.GraphOptimizationLevel.ORT_ENABLE_ALL
|
||||
sess_options.execution_mode = onnxruntime.ExecutionMode.ORT_SEQUENTIAL
|
||||
sess_options.intra_op_num_threads = 2
|
||||
|
||||
onnx_path = _find_first_existing_file(
|
||||
os.path.join(self.model_dir, "g2pW.onnx"),
|
||||
os.path.join(self.model_dir, "g2pw.onnx"),
|
||||
)
|
||||
|
||||
if "CUDAExecutionProvider" in onnxruntime.get_available_providers():
|
||||
self.session_g2pw = onnxruntime.InferenceSession(
|
||||
onnx_path,
|
||||
sess_options=sess_options,
|
||||
providers=["CUDAExecutionProvider", "CPUExecutionProvider"],
|
||||
)
|
||||
else:
|
||||
self.session_g2pw = onnxruntime.InferenceSession(
|
||||
onnx_path,
|
||||
sess_options=sess_options,
|
||||
providers=["CPUExecutionProvider"],
|
||||
)
|
||||
|
||||
def _predict(self, model_input: Dict[str, Any]) -> Tuple[List[str], List[float]]:
|
||||
return predict(session=self.session_g2pw, onnx_input=model_input, labels=self.labels)
|
||||
|
||||
@@ -45023,4 +45023,5 @@
|
||||
鼎铛玉石: ['ding3', 'cheng1', 'yu4', 'shi2']
|
||||
齿豁头童: ['chi3', 'huo1', 'tou2', 'tong2']
|
||||
牦牛: ['mao2', 'niu2']
|
||||
牦: ['mao2']
|
||||
牦: ['mao2']
|
||||
唑: ['zuo4']
|
||||
@@ -0,0 +1 @@
|
||||
13b2211c317c75794123ffdf7c2aea021c75b0c606ad61d7c8b05bb00b64fa21
|
||||
Binary file not shown.
@@ -655,11 +655,7 @@ class ToneSandhi:
|
||||
while i < len(seg):
|
||||
word, pos = seg[i]
|
||||
merged = False
|
||||
if (
|
||||
i - 1 >= 0
|
||||
and word == "一"
|
||||
and i + 1 < len(seg)
|
||||
):
|
||||
if i - 1 >= 0 and word == "一" and i + 1 < len(seg):
|
||||
last = new_seg[-1] if new_seg else seg[i - 1]
|
||||
if last[0] == seg[i + 1][0] and last[1] == "v" and seg[i + 1][1] == "v":
|
||||
combined = last[0] + "一" + seg[i + 1][0]
|
||||
|
||||
@@ -256,6 +256,24 @@ def replace_to_range(match) -> str:
|
||||
return result
|
||||
|
||||
|
||||
RE_VERSION_NUM = re.compile(r"((\d+)(\.\d+)(\.\d+)?(\.\d+)+)")
|
||||
def replace_vrsion_num(match) -> str:
|
||||
"""
|
||||
Args:
|
||||
match (re.Match)
|
||||
Returns:
|
||||
str
|
||||
"""
|
||||
result = ""
|
||||
for c in match.group(1):
|
||||
if c == ".":
|
||||
result += "点"
|
||||
else:
|
||||
result += num2str(c)
|
||||
return result
|
||||
|
||||
|
||||
|
||||
def _get_value(value_string: str, use_zero: bool = True) -> List[str]:
|
||||
stripped = value_string.lstrip("0")
|
||||
if len(stripped) == 0:
|
||||
@@ -308,7 +326,11 @@ def num2str(value_string: str) -> str:
|
||||
|
||||
result = verbalize_cardinal(integer)
|
||||
|
||||
decimal = decimal.rstrip("0")
|
||||
if decimal.endswith("0"):
|
||||
decimal = decimal.rstrip("0") + "0"
|
||||
else:
|
||||
decimal = decimal.rstrip("0")
|
||||
|
||||
if decimal:
|
||||
# '.22' is verbalized as '零点二二'
|
||||
# '3.20' is verbalized as '三点二
|
||||
|
||||
@@ -25,6 +25,7 @@ from .chronology import replace_time
|
||||
from .constants import F2H_ASCII_LETTERS
|
||||
from .constants import F2H_DIGITS
|
||||
from .constants import F2H_SPACE
|
||||
from .num import RE_VERSION_NUM
|
||||
from .num import RE_DECIMAL_NUM
|
||||
from .num import RE_DEFAULT_NUM
|
||||
from .num import RE_FRAC
|
||||
@@ -36,6 +37,7 @@ from .num import RE_RANGE
|
||||
from .num import RE_TO_RANGE
|
||||
from .num import RE_ASMD
|
||||
from .num import RE_POWER
|
||||
from .num import replace_vrsion_num
|
||||
from .num import replace_default_num
|
||||
from .num import replace_frac
|
||||
from .num import replace_negative_num
|
||||
@@ -158,6 +160,7 @@ class TextNormalizer:
|
||||
sentence = RE_RANGE.sub(replace_range, sentence)
|
||||
|
||||
sentence = RE_INTEGER.sub(replace_negative_num, sentence)
|
||||
sentence = RE_VERSION_NUM.sub(replace_vrsion_num, sentence)
|
||||
sentence = RE_DECIMAL_NUM.sub(replace_number, sentence)
|
||||
sentence = RE_POSITIVE_QUANTIFIERS.sub(replace_positive_quantifier, sentence)
|
||||
sentence = RE_DEFAULT_NUM.sub(replace_default_num, sentence)
|
||||
|
||||
+2
-2
@@ -283,7 +283,7 @@ def get_hparams_from_file(config_path):
|
||||
def check_git_hash(model_dir):
|
||||
source_dir = os.path.dirname(os.path.realpath(__file__))
|
||||
if not os.path.exists(os.path.join(source_dir, ".git")):
|
||||
logger.warn(
|
||||
logger.warning(
|
||||
"{} is not a git repository, therefore hash value comparison will be ignored.".format(
|
||||
source_dir,
|
||||
)
|
||||
@@ -296,7 +296,7 @@ def check_git_hash(model_dir):
|
||||
if os.path.exists(path):
|
||||
saved_hash = open(path).read()
|
||||
if saved_hash != cur_hash:
|
||||
logger.warn(
|
||||
logger.warning(
|
||||
"git hash values are different. {}(saved) != {}(current)".format(
|
||||
saved_hash[:8],
|
||||
cur_hash[:8],
|
||||
|
||||
@@ -9,10 +9,17 @@ A Powerful Few-shot Voice Conversion and Text-to-Speech WebUI.<br><br>
|
||||
|
||||
<!-- img src="https://counter.seku.su/cmoe?name=gptsovits&theme=r34" /><br> -->
|
||||
|
||||
[](https://colab.research.google.com/github/RVC-Boss/GPT-SoVITS/blob/main/colab_webui.ipynb)
|
||||
[](https://github.com/RVC-Boss/GPT-SoVITS/blob/main/LICENSE)
|
||||
[](https://huggingface.co/spaces/lj1995/GPT-SoVITS-v2)
|
||||
[](https://discord.gg/dnrgs5GHfG)
|
||||
[](https://www.python.org)
|
||||
[](https://github.com/RVC-Boss/gpt-sovits/releases)
|
||||
|
||||
[](https://colab.research.google.com/github/RVC-Boss/GPT-SoVITS/blob/main/Colab-WebUI.ipynb)
|
||||
[](https://lj1995-gpt-sovits-proplus.hf.space/)
|
||||
[](https://hub.docker.com/r/xxxxrt666/gpt-sovits)
|
||||
|
||||
[](https://www.yuque.com/baicaigongchang1145haoyuangong/ib3g1e)
|
||||
[](https://rentry.co/GPT-SoVITS-guide#/)
|
||||
[](https://github.com/RVC-Boss/GPT-SoVITS/blob/main/docs/en/Changelog_EN.md)
|
||||
[](https://github.com/RVC-Boss/GPT-SoVITS/blob/main/LICENSE)
|
||||
|
||||
**English** | [**中文简体**](./docs/cn/README.md) | [**日本語**](./docs/ja/README.md) | [**한국어**](./docs/ko/README.md) | [**Türkçe**](./docs/tr/README.md)
|
||||
|
||||
@@ -28,7 +35,7 @@ A Powerful Few-shot Voice Conversion and Text-to-Speech WebUI.<br><br>
|
||||
|
||||
3. **Cross-lingual Support:** Inference in languages different from the training dataset, currently supporting English, Japanese, Korean, Cantonese and Chinese.
|
||||
|
||||
4. **WebUI Tools:** Integrated tools include voice accompaniment separation, automatic training set segmentation, Chinese ASR, and text labeling, assisting beginners in creating training datasets and GPT/SoVITS models.
|
||||
4. **WebUI Tools:** Integrated tools include voice accompaniment separation, automatic training set segmentation, multilingual ASR with [Fun-ASR-Nano](https://github.com/FunAudioLLM/Fun-ASR), [SenseVoice](https://github.com/FunAudioLLM/SenseVoice), and classic [FunASR](https://github.com/modelscope/FunASR), plus text labeling, assisting beginners in creating training datasets and GPT/SoVITS models.
|
||||
|
||||
**Check out our [demo video](https://www.bilibili.com/video/BV12g4y1m7Uw) here!**
|
||||
|
||||
@@ -36,6 +43,13 @@ Unseen speakers few-shot fine-tuning demo:
|
||||
|
||||
https://github.com/RVC-Boss/GPT-SoVITS/assets/129054828/05bee1fa-bdd8-4d85-9350-80c060ab47fb
|
||||
|
||||
**RTF(inference speed) of GPT-SoVITS v2 ProPlus**:
|
||||
0.028 tested in 4060Ti, 0.014 tested in 4090 (1400words~=4min, inference time is 3.36s), 0.526 in M4 CPU. You can test our [huggingface demo](https://lj1995-gpt-sovits-proplus.hf.space/) (half H200) to experience high-speed inference .
|
||||
|
||||
请不要尬黑GPT-SoVITS推理速度慢,谢谢!
|
||||
|
||||
CPU-Optimized Inference Version:https://github.com/baicai-1145/GPT-SoVITS-CPUFast
|
||||
|
||||
**User guide: [简体中文](https://www.yuque.com/baicaigongchang1145haoyuangong/ib3g1e) | [English](https://rentry.co/GPT-SoVITS-guide#/)**
|
||||
|
||||
## Installation
|
||||
@@ -60,6 +74,14 @@ If you are a Windows user (tested with win>=10), you can [download the integrate
|
||||
|
||||
**Users in China can [download the package here](https://www.yuque.com/baicaigongchang1145haoyuangong/ib3g1e/dkxgpiy9zb96hob4#KTvnO).**
|
||||
|
||||
Install the program by running the following commands:
|
||||
|
||||
```pwsh
|
||||
conda create -n GPTSoVits python=3.10
|
||||
conda activate GPTSoVits
|
||||
pwsh -F install.ps1 --Device <CU126|CU128|CPU> --Source <HF|HF-Mirror|ModelScope> [--DownloadUVR5]
|
||||
```
|
||||
|
||||
### Linux
|
||||
|
||||
```bash
|
||||
@@ -128,8 +150,9 @@ Due to rapid development in the codebase and a slower Docker image release cycle
|
||||
|
||||
- Check [Docker Hub](https://hub.docker.com/r/xxxxrt666/gpt-sovits) for the latest available image tags
|
||||
- Choose an appropriate image tag for your environment
|
||||
- `Lite` means the Docker image does not include ASR models and UVR5 models. You can manually download the UVR5 models, while the program will automatically download the ASR models as needed
|
||||
- `Lite` means the Docker image **does not include** ASR models and UVR5 models. You can manually download the UVR5 models, while the program will automatically download the ASR models as needed
|
||||
- The appropriate architecture image (amd64/arm64) will be automatically pulled during Docker Compose
|
||||
- Docker Compose will mount **all files** in the current directory. Please switch to the project root directory and **pull the latest code** before using the Docker image
|
||||
- Optionally, build the image locally using the provided Dockerfile for the most up-to-date changes
|
||||
|
||||
#### Environment Variables
|
||||
@@ -185,7 +208,7 @@ docker exec -it <GPT-SoVITS-CU126-Lite|GPT-SoVITS-CU128-Lite|GPT-SoVITS-CU126|GP
|
||||
|
||||
- The suggestion is to **directly specify the model type** in the model name and configuration file name, such as `mel_mand_roformer`, `bs_roformer`. If not specified, the features will be compared from the configuration file to determine which type of model it is. For example, the model `bs_roformer_ep_368_sdr_12.9628.ckpt` and its corresponding configuration file `bs_roformer_ep_368_sdr_12.9628.yaml` are a pair, `kim_mel_band_roformer.ckpt` and `kim_mel_band_roformer.yaml` are also a pair.
|
||||
|
||||
4. For Chinese ASR (additionally), download models from [Damo ASR Model](https://modelscope.cn/models/damo/speech_paraformer-large_asr_nat-zh-cn-16k-common-vocab8404-pytorch/files), [Damo VAD Model](https://modelscope.cn/models/damo/speech_fsmn_vad_zh-cn-16k-common-pytorch/files), and [Damo Punc Model](https://modelscope.cn/models/damo/punc_ct-transformer_zh-cn-common-vocab272727-pytorch/files) and place them in `tools/asr/models`.
|
||||
4. FunASR models are downloaded automatically on first use. The WebUI offers [Fun-ASR-Nano](https://github.com/FunAudioLLM/Fun-ASR) for multilingual and dialect ASR, [SenseVoice](https://github.com/FunAudioLLM/SenseVoice) for fast transcription, and classic Paraformer/UniASR through [FunASR](https://github.com/modelscope/FunASR) for Chinese and Cantonese. To preinstall the classic Chinese models for offline use, download the [ASR model](https://modelscope.cn/models/iic/speech_paraformer-large_asr_nat-zh-cn-16k-common-vocab8404-pytorch/files), [VAD model](https://modelscope.cn/models/iic/speech_fsmn_vad_zh-cn-16k-common-pytorch/files), and [punctuation model](https://modelscope.cn/models/iic/punc_ct-transformer_zh-cn-common-vocab272727-pytorch/files) into `tools/asr/models`.
|
||||
|
||||
5. For English or Japanese ASR (additionally), download models from [Faster Whisper Large V3](https://huggingface.co/Systran/faster-whisper-large-v3) and place them in `tools/asr/models`. Also, [other models](https://huggingface.co/Systran) may have the similar effect with smaller disk footprint.
|
||||
|
||||
@@ -326,14 +349,14 @@ Use v4 from v1/v2/v3 environment:
|
||||
|
||||
2. Clone the latest codes from github.
|
||||
|
||||
3. Download v4 pretrained models (gsv-v4-pretrained/s2v4.ckpt, and gsv-v4-pretrained/vocoder.pth) from [huggingface](https://huggingface.co/lj1995/GPT-SoVITS/tree/main) and put them into `GPT_SoVITS/pretrained_models`.
|
||||
3. Download v4 pretrained models (gsv-v4-pretrained/s2v4.pth, and gsv-v4-pretrained/vocoder.pth) from [huggingface](https://huggingface.co/lj1995/GPT-SoVITS/tree/main) and put them into `GPT_SoVITS/pretrained_models`.
|
||||
|
||||
## V2Pro Release Notes
|
||||
|
||||
New Features:
|
||||
|
||||
1. Slightly higher VRAM usage than v2, surpassing v4's performance, with v2's hardware cost and speed.
|
||||
[more details](https://github.com/RVC-Boss/GPT-SoVITS/wiki/GPT%E2%80%90SoVITS%E2%80%90features-(%E5%90%84%E7%89%88%E6%9C%AC%E7%89%B9%E6%80%A7))
|
||||
[more details](<https://github.com/RVC-Boss/GPT-SoVITS/wiki/GPT%E2%80%90SoVITS%E2%80%90features-(%E5%90%84%E7%89%88%E6%9C%AC%E7%89%B9%E6%80%A7)>)
|
||||
|
||||
2.v1/v2 and the v2Pro series share the same characteristics, while v3/v4 have similar features. For training sets with average audio quality, v1/v2/v2Pro can deliver decent results, but v3/v4 cannot. Additionally, the synthesized tone and timebre of v3/v4 lean more toward the reference audio rather than the overall training set.
|
||||
|
||||
@@ -390,13 +413,13 @@ python audio_slicer.py \
|
||||
--hop_size <step_size_for_computing_volume_curve>
|
||||
```
|
||||
|
||||
This is how dataset ASR processing is done using the command line(Only Chinese)
|
||||
Run dataset ASR with FunASR from the command line. Fun-ASR-Nano is the default for Chinese, English, Japanese, Korean, and automatic language detection; Cantonese keeps the classic FunASR backend.
|
||||
|
||||
```bash
|
||||
python tools/asr/funasr_asr.py -i <input> -o <output>
|
||||
python tools/asr/funasr_asr.py -i <input> -o <output> -l zh
|
||||
```
|
||||
|
||||
ASR processing is performed through Faster_Whisper(ASR marking except Chinese)
|
||||
Faster Whisper is also available as an ASR backend.
|
||||
|
||||
(No progress bars, GPU performance may cause time delays)
|
||||
|
||||
@@ -445,7 +468,9 @@ Special thanks to the following projects and contributors:
|
||||
- [FFmpeg](https://github.com/FFmpeg/FFmpeg)
|
||||
- [gradio](https://github.com/gradio-app/gradio)
|
||||
- [faster-whisper](https://github.com/SYSTRAN/faster-whisper)
|
||||
- [FunASR](https://github.com/alibaba-damo-academy/FunASR)
|
||||
- [FunASR](https://github.com/modelscope/FunASR)
|
||||
- [Fun-ASR](https://github.com/FunAudioLLM/Fun-ASR)
|
||||
- [SenseVoice](https://github.com/FunAudioLLM/SenseVoice)
|
||||
- [AP-BWE](https://github.com/yxlu-0102/AP-BWE)
|
||||
|
||||
Thankful to @Naozumi520 for providing the Cantonese training set and for the guidance on Cantonese-related knowledge.
|
||||
|
||||
@@ -163,7 +163,7 @@ from transformers import AutoModelForMaskedLM, AutoTokenizer
|
||||
import numpy as np
|
||||
from feature_extractor import cnhubert
|
||||
from io import BytesIO
|
||||
from module.models import SynthesizerTrn, SynthesizerTrnV3
|
||||
from module.models import Generator, SynthesizerTrn, SynthesizerTrnV3
|
||||
from peft import LoraConfig, get_peft_model
|
||||
from AR.models.t2s_lightning_module import Text2SemanticLightningModule
|
||||
from text import cleaned_text_to_sequence
|
||||
@@ -198,31 +198,103 @@ def is_full(*items): # 任意一项为空返回False
|
||||
return True
|
||||
|
||||
|
||||
def init_bigvgan():
|
||||
bigvgan_model = hifigan_model = sv_cn_model = None
|
||||
|
||||
|
||||
def clean_hifigan_model():
|
||||
global hifigan_model
|
||||
if hifigan_model:
|
||||
hifigan_model = hifigan_model.cpu()
|
||||
hifigan_model = None
|
||||
try:
|
||||
torch.cuda.empty_cache()
|
||||
except:
|
||||
pass
|
||||
|
||||
|
||||
def clean_bigvgan_model():
|
||||
global bigvgan_model
|
||||
if bigvgan_model:
|
||||
bigvgan_model = bigvgan_model.cpu()
|
||||
bigvgan_model = None
|
||||
try:
|
||||
torch.cuda.empty_cache()
|
||||
except:
|
||||
pass
|
||||
|
||||
|
||||
def clean_sv_cn_model():
|
||||
global sv_cn_model
|
||||
if sv_cn_model:
|
||||
sv_cn_model.embedding_model = sv_cn_model.embedding_model.cpu()
|
||||
sv_cn_model = None
|
||||
try:
|
||||
torch.cuda.empty_cache()
|
||||
except:
|
||||
pass
|
||||
|
||||
|
||||
def init_bigvgan():
|
||||
global bigvgan_model, hifigan_model, sv_cn_model
|
||||
from BigVGAN import bigvgan
|
||||
|
||||
bigvgan_model = bigvgan.BigVGAN.from_pretrained(
|
||||
"%s/GPT_SoVITS/pretrained_models/models--nvidia--bigvgan_v2_24khz_100band_256x" % (now_dir,),
|
||||
use_cuda_kernel=False,
|
||||
) # if True, RuntimeError: Ninja is required to load C++ extensions
|
||||
)
|
||||
# remove weight norm in the model and set to eval mode
|
||||
bigvgan_model.remove_weight_norm()
|
||||
bigvgan_model = bigvgan_model.eval()
|
||||
|
||||
if is_half == True:
|
||||
bigvgan_model = bigvgan_model.half().to(device)
|
||||
else:
|
||||
bigvgan_model = bigvgan_model.to(device)
|
||||
|
||||
|
||||
def init_hifigan():
|
||||
global hifigan_model, bigvgan_model, sv_cn_model
|
||||
hifigan_model = Generator(
|
||||
initial_channel=100,
|
||||
resblock="1",
|
||||
resblock_kernel_sizes=[3, 7, 11],
|
||||
resblock_dilation_sizes=[[1, 3, 5], [1, 3, 5], [1, 3, 5]],
|
||||
upsample_rates=[10, 6, 2, 2, 2],
|
||||
upsample_initial_channel=512,
|
||||
upsample_kernel_sizes=[20, 12, 4, 4, 4],
|
||||
gin_channels=0,
|
||||
is_bias=True,
|
||||
)
|
||||
hifigan_model.eval()
|
||||
hifigan_model.remove_weight_norm()
|
||||
state_dict_g = torch.load(
|
||||
"%s/GPT_SoVITS/pretrained_models/gsv-v4-pretrained/vocoder.pth" % (now_dir,),
|
||||
map_location="cpu",
|
||||
weights_only=False,
|
||||
)
|
||||
print("loading vocoder", hifigan_model.load_state_dict(state_dict_g))
|
||||
if is_half == True:
|
||||
hifigan_model = hifigan_model.half().to(device)
|
||||
else:
|
||||
hifigan_model = hifigan_model.to(device)
|
||||
|
||||
|
||||
from sv import SV
|
||||
|
||||
|
||||
def init_sv_cn():
|
||||
global hifigan_model, bigvgan_model, sv_cn_model
|
||||
sv_cn_model = SV(device, is_half)
|
||||
|
||||
|
||||
resample_transform_dict = {}
|
||||
|
||||
|
||||
def resample(audio_tensor, sr0):
|
||||
def resample(audio_tensor, sr0, sr1, device):
|
||||
global resample_transform_dict
|
||||
if sr0 not in resample_transform_dict:
|
||||
resample_transform_dict[sr0] = torchaudio.transforms.Resample(sr0, 24000).to(device)
|
||||
return resample_transform_dict[sr0](audio_tensor)
|
||||
key = "%s-%s-%s" % (sr0, sr1, str(device))
|
||||
if key not in resample_transform_dict:
|
||||
resample_transform_dict[key] = torchaudio.transforms.Resample(sr0, sr1).to(device)
|
||||
return resample_transform_dict[key](audio_tensor)
|
||||
|
||||
|
||||
from module.mel_processing import mel_spectrogram_torch
|
||||
@@ -252,6 +324,19 @@ mel_fn = lambda x: mel_spectrogram_torch(
|
||||
"center": False,
|
||||
},
|
||||
)
|
||||
mel_fn_v4 = lambda x: mel_spectrogram_torch(
|
||||
x,
|
||||
**{
|
||||
"n_fft": 1280,
|
||||
"win_size": 1280,
|
||||
"hop_size": 320,
|
||||
"num_mels": 100,
|
||||
"sampling_rate": 32000,
|
||||
"fmin": 0,
|
||||
"fmax": None,
|
||||
"center": False,
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
sr_model = None
|
||||
@@ -293,12 +378,19 @@ from process_ckpt import get_sovits_version_from_path_fast, load_sovits_new
|
||||
|
||||
|
||||
def get_sovits_weights(sovits_path):
|
||||
path_sovits_v3 = "GPT_SoVITS/pretrained_models/s2Gv3.pth"
|
||||
from config import pretrained_sovits_name
|
||||
|
||||
path_sovits_v3 = pretrained_sovits_name["v3"]
|
||||
path_sovits_v4 = pretrained_sovits_name["v4"]
|
||||
is_exist_s2gv3 = os.path.exists(path_sovits_v3)
|
||||
is_exist_s2gv4 = os.path.exists(path_sovits_v4)
|
||||
|
||||
version, model_version, if_lora_v3 = get_sovits_version_from_path_fast(sovits_path)
|
||||
if if_lora_v3 == True and is_exist_s2gv3 == False:
|
||||
logger.info("SoVITS V3 底模缺失,无法加载相应 LoRA 权重")
|
||||
is_exist = is_exist_s2gv3 if model_version == "v3" else is_exist_s2gv4
|
||||
path_sovits = path_sovits_v3 if model_version == "v3" else path_sovits_v4
|
||||
|
||||
if if_lora_v3 == True and is_exist == False:
|
||||
logger.info("SoVITS %s 底模缺失,无法加载相应 LoRA 权重" % model_version)
|
||||
|
||||
dict_s2 = load_sovits_new(sovits_path)
|
||||
hps = dict_s2["config"]
|
||||
@@ -311,11 +403,13 @@ def get_sovits_weights(sovits_path):
|
||||
else:
|
||||
hps.model.version = "v2"
|
||||
|
||||
if model_version == "v3":
|
||||
hps.model.version = "v3"
|
||||
|
||||
model_params_dict = vars(hps.model)
|
||||
if model_version != "v3":
|
||||
if model_version not in {"v3", "v4"}:
|
||||
if "Pro" in model_version:
|
||||
hps.model.version = model_version
|
||||
if sv_cn_model == None:
|
||||
init_sv_cn()
|
||||
|
||||
vq_model = SynthesizerTrn(
|
||||
hps.data.filter_length // 2 + 1,
|
||||
hps.train.segment_size // hps.data.hop_length,
|
||||
@@ -323,13 +417,18 @@ def get_sovits_weights(sovits_path):
|
||||
**model_params_dict,
|
||||
)
|
||||
else:
|
||||
hps.model.version = model_version
|
||||
vq_model = SynthesizerTrnV3(
|
||||
hps.data.filter_length // 2 + 1,
|
||||
hps.train.segment_size // hps.data.hop_length,
|
||||
n_speakers=hps.data.n_speakers,
|
||||
**model_params_dict,
|
||||
)
|
||||
init_bigvgan()
|
||||
if model_version == "v3":
|
||||
init_bigvgan()
|
||||
if model_version == "v4":
|
||||
init_hifigan()
|
||||
|
||||
model_version = hps.model.version
|
||||
logger.info(f"模型版本: {model_version}")
|
||||
if "pretrained" not in sovits_path:
|
||||
@@ -345,7 +444,8 @@ def get_sovits_weights(sovits_path):
|
||||
if if_lora_v3 == False:
|
||||
vq_model.load_state_dict(dict_s2["weight"], strict=False)
|
||||
else:
|
||||
vq_model.load_state_dict(load_sovits_new(path_sovits_v3)["weight"], strict=False)
|
||||
path_sovits = path_sovits_v3 if model_version == "v3" else path_sovits_v4
|
||||
vq_model.load_state_dict(load_sovits_new(path_sovits)["weight"], strict=False)
|
||||
lora_rank = dict_s2["lora_rank"]
|
||||
lora_config = LoraConfig(
|
||||
target_modules=["to_k", "to_q", "to_v", "to_out.0"],
|
||||
@@ -442,62 +542,65 @@ from text import chinese
|
||||
|
||||
|
||||
def get_phones_and_bert(text, language, version, final=False):
|
||||
if language in {"en", "all_zh", "all_ja", "all_ko", "all_yue"}:
|
||||
formattext = text
|
||||
while " " in formattext:
|
||||
formattext = formattext.replace(" ", " ")
|
||||
if language == "all_zh":
|
||||
if re.search(r"[A-Za-z]", formattext):
|
||||
formattext = re.sub(r"[a-z]", lambda x: x.group(0).upper(), formattext)
|
||||
formattext = chinese.mix_text_normalize(formattext)
|
||||
return get_phones_and_bert(formattext, "zh", version)
|
||||
text = re.sub(r' {2,}', ' ', text)
|
||||
textlist = []
|
||||
langlist = []
|
||||
if language == "all_zh":
|
||||
for tmp in LangSegmenter.getTexts(text,"zh"):
|
||||
langlist.append(tmp["lang"])
|
||||
textlist.append(tmp["text"])
|
||||
elif language == "all_yue":
|
||||
for tmp in LangSegmenter.getTexts(text,"zh"):
|
||||
if tmp["lang"] == "zh":
|
||||
tmp["lang"] = "yue"
|
||||
langlist.append(tmp["lang"])
|
||||
textlist.append(tmp["text"])
|
||||
elif language == "all_ja":
|
||||
for tmp in LangSegmenter.getTexts(text,"ja"):
|
||||
langlist.append(tmp["lang"])
|
||||
textlist.append(tmp["text"])
|
||||
elif language == "all_ko":
|
||||
for tmp in LangSegmenter.getTexts(text,"ko"):
|
||||
langlist.append(tmp["lang"])
|
||||
textlist.append(tmp["text"])
|
||||
elif language == "en":
|
||||
langlist.append("en")
|
||||
textlist.append(text)
|
||||
elif language == "auto":
|
||||
for tmp in LangSegmenter.getTexts(text):
|
||||
langlist.append(tmp["lang"])
|
||||
textlist.append(tmp["text"])
|
||||
elif language == "auto_yue":
|
||||
for tmp in LangSegmenter.getTexts(text):
|
||||
if tmp["lang"] == "zh":
|
||||
tmp["lang"] = "yue"
|
||||
langlist.append(tmp["lang"])
|
||||
textlist.append(tmp["text"])
|
||||
else:
|
||||
for tmp in LangSegmenter.getTexts(text):
|
||||
if langlist:
|
||||
if (tmp["lang"] == "en" and langlist[-1] == "en") or (tmp["lang"] != "en" and langlist[-1] != "en"):
|
||||
textlist[-1] += tmp["text"]
|
||||
continue
|
||||
if tmp["lang"] == "en":
|
||||
langlist.append(tmp["lang"])
|
||||
else:
|
||||
phones, word2ph, norm_text = clean_text_inf(formattext, language, version)
|
||||
bert = get_bert_feature(norm_text, word2ph).to(device)
|
||||
elif language == "all_yue" and re.search(r"[A-Za-z]", formattext):
|
||||
formattext = re.sub(r"[a-z]", lambda x: x.group(0).upper(), formattext)
|
||||
formattext = chinese.mix_text_normalize(formattext)
|
||||
return get_phones_and_bert(formattext, "yue", version)
|
||||
else:
|
||||
phones, word2ph, norm_text = clean_text_inf(formattext, language, version)
|
||||
bert = torch.zeros(
|
||||
(1024, len(phones)),
|
||||
dtype=torch.float16 if is_half == True else torch.float32,
|
||||
).to(device)
|
||||
elif language in {"zh", "ja", "ko", "yue", "auto", "auto_yue"}:
|
||||
textlist = []
|
||||
langlist = []
|
||||
if language == "auto":
|
||||
for tmp in LangSegmenter.getTexts(text):
|
||||
langlist.append(tmp["lang"])
|
||||
textlist.append(tmp["text"])
|
||||
elif language == "auto_yue":
|
||||
for tmp in LangSegmenter.getTexts(text):
|
||||
if tmp["lang"] == "zh":
|
||||
tmp["lang"] = "yue"
|
||||
langlist.append(tmp["lang"])
|
||||
textlist.append(tmp["text"])
|
||||
else:
|
||||
for tmp in LangSegmenter.getTexts(text):
|
||||
if tmp["lang"] == "en":
|
||||
langlist.append(tmp["lang"])
|
||||
else:
|
||||
# 因无法区别中日韩文汉字,以用户输入为准
|
||||
langlist.append(language)
|
||||
textlist.append(tmp["text"])
|
||||
phones_list = []
|
||||
bert_list = []
|
||||
norm_text_list = []
|
||||
for i in range(len(textlist)):
|
||||
lang = langlist[i]
|
||||
phones, word2ph, norm_text = clean_text_inf(textlist[i], lang, version)
|
||||
bert = get_bert_inf(phones, word2ph, norm_text, lang)
|
||||
phones_list.append(phones)
|
||||
norm_text_list.append(norm_text)
|
||||
bert_list.append(bert)
|
||||
bert = torch.cat(bert_list, dim=1)
|
||||
phones = sum(phones_list, [])
|
||||
norm_text = "".join(norm_text_list)
|
||||
# 因无法区别中日韩文汉字,以用户输入为准
|
||||
langlist.append(language)
|
||||
textlist.append(tmp["text"])
|
||||
phones_list = []
|
||||
bert_list = []
|
||||
norm_text_list = []
|
||||
for i in range(len(textlist)):
|
||||
lang = langlist[i]
|
||||
phones, word2ph, norm_text = clean_text_inf(textlist[i], lang, version)
|
||||
bert = get_bert_inf(phones, word2ph, norm_text, lang)
|
||||
phones_list.append(phones)
|
||||
norm_text_list.append(norm_text)
|
||||
bert_list.append(bert)
|
||||
bert = torch.cat(bert_list, dim=1)
|
||||
phones = sum(phones_list, [])
|
||||
norm_text = "".join(norm_text_list)
|
||||
|
||||
if not final and len(phones) < 6:
|
||||
return get_phones_and_bert("." + text, language, version, final=True)
|
||||
@@ -533,23 +636,34 @@ class DictToAttrRecursive(dict):
|
||||
raise AttributeError(f"Attribute {item} not found")
|
||||
|
||||
|
||||
def get_spepc(hps, filename):
|
||||
audio, _ = librosa.load(filename, sr=int(hps.data.sampling_rate))
|
||||
audio = torch.FloatTensor(audio)
|
||||
def get_spepc(hps, filename, dtype, device, is_v2pro=False):
|
||||
sr1 = int(hps.data.sampling_rate)
|
||||
audio, sr0 = torchaudio.load(filename)
|
||||
if sr0 != sr1:
|
||||
audio = audio.to(device)
|
||||
if audio.shape[0] == 2:
|
||||
audio = audio.mean(0).unsqueeze(0)
|
||||
audio = resample(audio, sr0, sr1, device)
|
||||
else:
|
||||
audio = audio.to(device)
|
||||
if audio.shape[0] == 2:
|
||||
audio = audio.mean(0).unsqueeze(0)
|
||||
|
||||
maxx = audio.abs().max()
|
||||
if maxx > 1:
|
||||
audio /= min(2, maxx)
|
||||
audio_norm = audio
|
||||
audio_norm = audio_norm.unsqueeze(0)
|
||||
spec = spectrogram_torch(
|
||||
audio_norm,
|
||||
audio,
|
||||
hps.data.filter_length,
|
||||
hps.data.sampling_rate,
|
||||
hps.data.hop_length,
|
||||
hps.data.win_length,
|
||||
center=False,
|
||||
)
|
||||
return spec
|
||||
spec = spec.to(dtype)
|
||||
if is_v2pro == True:
|
||||
audio = resample(audio, sr1, 16000, device).to(dtype)
|
||||
return spec, audio
|
||||
|
||||
|
||||
def pack_audio(audio_bytes, data, rate):
|
||||
@@ -736,6 +850,16 @@ def get_tts_wav(
|
||||
t2s_model = infer_gpt.t2s_model
|
||||
max_sec = infer_gpt.max_sec
|
||||
|
||||
if version == "v3":
|
||||
if sample_steps not in [4, 8, 16, 32, 64, 128]:
|
||||
sample_steps = 32
|
||||
elif version == "v4":
|
||||
if sample_steps not in [4, 8, 16, 32]:
|
||||
sample_steps = 8
|
||||
|
||||
if if_sr and version != "v3":
|
||||
if_sr = False
|
||||
|
||||
t0 = ttime()
|
||||
prompt_text = prompt_text.strip("\n")
|
||||
if prompt_text[-1] not in splits:
|
||||
@@ -759,19 +883,29 @@ def get_tts_wav(
|
||||
prompt_semantic = codes[0, 0]
|
||||
prompt = prompt_semantic.unsqueeze(0).to(device)
|
||||
|
||||
if version != "v3":
|
||||
is_v2pro = version in {"v2Pro", "v2ProPlus"}
|
||||
if version not in {"v3", "v4"}:
|
||||
refers = []
|
||||
if is_v2pro:
|
||||
sv_emb = []
|
||||
if sv_cn_model == None:
|
||||
init_sv_cn()
|
||||
if inp_refs:
|
||||
for path in inp_refs:
|
||||
try:
|
||||
refer = get_spepc(hps, path).to(dtype).to(device)
|
||||
try: #####这里加上提取sv的逻辑,要么一堆sv一堆refer,要么单个sv单个refer
|
||||
refer, audio_tensor = get_spepc(hps, path.name, dtype, device, is_v2pro)
|
||||
refers.append(refer)
|
||||
if is_v2pro:
|
||||
sv_emb.append(sv_cn_model.compute_embedding3(audio_tensor))
|
||||
except Exception as e:
|
||||
logger.error(e)
|
||||
if len(refers) == 0:
|
||||
refers = [get_spepc(hps, ref_wav_path).to(dtype).to(device)]
|
||||
refers, audio_tensor = get_spepc(hps, ref_wav_path, dtype, device, is_v2pro)
|
||||
refers = [refers]
|
||||
if is_v2pro:
|
||||
sv_emb = [sv_cn_model.compute_embedding3(audio_tensor)]
|
||||
else:
|
||||
refer = get_spepc(hps, ref_wav_path).to(device).to(dtype)
|
||||
refer, audio_tensor = get_spepc(hps, ref_wav_path, dtype, device)
|
||||
|
||||
t1 = ttime()
|
||||
# os.environ['version'] = version
|
||||
@@ -811,41 +945,56 @@ def get_tts_wav(
|
||||
pred_semantic = pred_semantic[:, -idx:].unsqueeze(0)
|
||||
t3 = ttime()
|
||||
|
||||
if version != "v3":
|
||||
audio = (
|
||||
vq_model.decode(pred_semantic, torch.LongTensor(phones2).to(device).unsqueeze(0), refers, speed=speed)
|
||||
.detach()
|
||||
.cpu()
|
||||
.numpy()[0, 0]
|
||||
) ###试试重建不带上prompt部分
|
||||
if version not in {"v3", "v4"}:
|
||||
if is_v2pro:
|
||||
audio = (
|
||||
vq_model.decode(
|
||||
pred_semantic,
|
||||
torch.LongTensor(phones2).to(device).unsqueeze(0),
|
||||
refers,
|
||||
speed=speed,
|
||||
sv_emb=sv_emb,
|
||||
)
|
||||
.detach()
|
||||
.cpu()
|
||||
.numpy()[0, 0]
|
||||
)
|
||||
else:
|
||||
audio = (
|
||||
vq_model.decode(
|
||||
pred_semantic, torch.LongTensor(phones2).to(device).unsqueeze(0), refers, speed=speed
|
||||
)
|
||||
.detach()
|
||||
.cpu()
|
||||
.numpy()[0, 0]
|
||||
)
|
||||
else:
|
||||
phoneme_ids0 = torch.LongTensor(phones1).to(device).unsqueeze(0)
|
||||
phoneme_ids1 = torch.LongTensor(phones2).to(device).unsqueeze(0)
|
||||
# print(11111111, phoneme_ids0, phoneme_ids1)
|
||||
|
||||
fea_ref, ge = vq_model.decode_encp(prompt.unsqueeze(0), phoneme_ids0, refer)
|
||||
ref_audio, sr = torchaudio.load(ref_wav_path)
|
||||
ref_audio = ref_audio.to(device).float()
|
||||
if ref_audio.shape[0] == 2:
|
||||
ref_audio = ref_audio.mean(0).unsqueeze(0)
|
||||
if sr != 24000:
|
||||
ref_audio = resample(ref_audio, sr)
|
||||
# print("ref_audio",ref_audio.abs().mean())
|
||||
mel2 = mel_fn(ref_audio)
|
||||
|
||||
tgt_sr = 24000 if version == "v3" else 32000
|
||||
if sr != tgt_sr:
|
||||
ref_audio = resample(ref_audio, sr, tgt_sr, device)
|
||||
mel2 = mel_fn(ref_audio) if version == "v3" else mel_fn_v4(ref_audio)
|
||||
mel2 = norm_spec(mel2)
|
||||
T_min = min(mel2.shape[2], fea_ref.shape[2])
|
||||
mel2 = mel2[:, :, :T_min]
|
||||
fea_ref = fea_ref[:, :, :T_min]
|
||||
if T_min > 468:
|
||||
mel2 = mel2[:, :, -468:]
|
||||
fea_ref = fea_ref[:, :, -468:]
|
||||
T_min = 468
|
||||
chunk_len = 934 - T_min
|
||||
# print("fea_ref",fea_ref,fea_ref.shape)
|
||||
# print("mel2",mel2)
|
||||
Tref = 468 if version == "v3" else 500
|
||||
Tchunk = 934 if version == "v3" else 1000
|
||||
if T_min > Tref:
|
||||
mel2 = mel2[:, :, -Tref:]
|
||||
fea_ref = fea_ref[:, :, -Tref:]
|
||||
T_min = Tref
|
||||
chunk_len = Tchunk - T_min
|
||||
mel2 = mel2.to(dtype)
|
||||
fea_todo, ge = vq_model.decode_encp(pred_semantic, phoneme_ids1, refer, ge, speed)
|
||||
# print("fea_todo",fea_todo)
|
||||
# print("ge",ge.abs().mean())
|
||||
cfm_resss = []
|
||||
idx = 0
|
||||
while 1:
|
||||
@@ -854,22 +1003,24 @@ def get_tts_wav(
|
||||
break
|
||||
idx += chunk_len
|
||||
fea = torch.cat([fea_ref, fea_todo_chunk], 2).transpose(2, 1)
|
||||
# set_seed(123)
|
||||
cfm_res = vq_model.cfm.inference(
|
||||
fea, torch.LongTensor([fea.size(1)]).to(fea.device), mel2, sample_steps, inference_cfg_rate=0
|
||||
)
|
||||
cfm_res = cfm_res[:, :, mel2.shape[2] :]
|
||||
mel2 = cfm_res[:, :, -T_min:]
|
||||
# print("fea", fea)
|
||||
# print("mel2in", mel2)
|
||||
fea_ref = fea_todo_chunk[:, :, -T_min:]
|
||||
cfm_resss.append(cfm_res)
|
||||
cmf_res = torch.cat(cfm_resss, 2)
|
||||
cmf_res = denorm_spec(cmf_res)
|
||||
if bigvgan_model == None:
|
||||
init_bigvgan()
|
||||
cfm_res = torch.cat(cfm_resss, 2)
|
||||
cfm_res = denorm_spec(cfm_res)
|
||||
if version == "v3":
|
||||
if bigvgan_model == None:
|
||||
init_bigvgan()
|
||||
else: # v4
|
||||
if hifigan_model == None:
|
||||
init_hifigan()
|
||||
vocoder_model = bigvgan_model if version == "v3" else hifigan_model
|
||||
with torch.inference_mode():
|
||||
wav_gen = bigvgan_model(cmf_res)
|
||||
wav_gen = vocoder_model(cfm_res)
|
||||
audio = wav_gen[0][0].cpu().detach().numpy()
|
||||
|
||||
max_audio = np.abs(audio).max()
|
||||
@@ -880,7 +1031,13 @@ def get_tts_wav(
|
||||
audio_opt = np.concatenate(audio_opt, 0)
|
||||
t4 = ttime()
|
||||
|
||||
sr = hps.data.sampling_rate if version != "v3" else 24000
|
||||
if version in {"v1", "v2", "v2Pro", "v2ProPlus"}:
|
||||
sr = 32000
|
||||
elif version == "v3":
|
||||
sr = 24000
|
||||
else:
|
||||
sr = 48000 # v4
|
||||
|
||||
if if_sr and sr == 24000:
|
||||
audio_opt = torch.from_numpy(audio_opt).float().to(device)
|
||||
audio_opt, sr = audio_sr(audio_opt.unsqueeze(0), sr)
|
||||
@@ -900,8 +1057,12 @@ def get_tts_wav(
|
||||
|
||||
if not stream_mode == "normal":
|
||||
if media_type == "wav":
|
||||
sr = 48000 if if_sr else 24000
|
||||
sr = hps.data.sampling_rate if version != "v3" else sr
|
||||
if version in {"v1", "v2", "v2Pro", "v2ProPlus"}:
|
||||
sr = 32000
|
||||
elif version == "v3":
|
||||
sr = 48000 if if_sr else 24000
|
||||
else:
|
||||
sr = 48000 # v4
|
||||
audio_bytes = pack_wav(audio_bytes, sr)
|
||||
yield audio_bytes.getvalue()
|
||||
|
||||
@@ -966,9 +1127,6 @@ def handle(
|
||||
if not default_refer.is_ready():
|
||||
return JSONResponse({"code": 400, "message": "未指定参考音频且接口无预设"}, status_code=400)
|
||||
|
||||
if sample_steps not in [4, 8, 16, 32]:
|
||||
sample_steps = 32
|
||||
|
||||
if cut_punc == None:
|
||||
text = cut_text(text, default_cut_punc)
|
||||
else:
|
||||
@@ -1071,10 +1229,10 @@ default_refer = DefaultRefer(args.default_refer_path, args.default_refer_text, a
|
||||
# 模型路径检查
|
||||
if sovits_path == "":
|
||||
sovits_path = g_config.pretrained_sovits_path
|
||||
logger.warn(f"未指定SoVITS模型路径, fallback后当前值: {sovits_path}")
|
||||
logger.warning(f"未指定SoVITS模型路径, fallback后当前值: {sovits_path}")
|
||||
if gpt_path == "":
|
||||
gpt_path = g_config.pretrained_gpt_path
|
||||
logger.warn(f"未指定GPT模型路径, fallback后当前值: {gpt_path}")
|
||||
logger.warning(f"未指定GPT模型路径, fallback后当前值: {gpt_path}")
|
||||
|
||||
# 指定默认参考音频, 调用方 未提供/未给全 参考音频参数时使用
|
||||
if default_refer.path == "" or default_refer.text == "" or default_refer.language == "":
|
||||
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user