Repackage frontend as third-party dependency
This commit is contained in:
Owen Song
system
+152
@@ -0,0 +1,152 @@
|
||||
Apache License
|
||||
Version 2.0, January 2004
|
||||
http://www.apache.org/licenses/
|
||||
|
||||
TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
|
||||
|
||||
1. Definitions.
|
||||
|
||||
"License" shall mean the terms and conditions for use, reproduction,
|
||||
and distribution as defined by Sections 1 through 9 of this document.
|
||||
|
||||
"Licensor" shall mean the copyright owner or entity authorized by
|
||||
the copyright owner that is granting the License.
|
||||
|
||||
"Legal Entity" shall mean the union of the acting entity and all
|
||||
other entities that control, are controlled by, or are under common
|
||||
control with that entity. For the purposes of this definition,
|
||||
"control" means (i) the power, direct or indirect, to cause the
|
||||
direction or management of such entity, whether by contract or
|
||||
otherwise, or (ii) ownership of fifty percent (50%) or more of the
|
||||
outstanding shares, or (iii) beneficial ownership of such entity.
|
||||
|
||||
"You" (or "Your") shall mean an individual or Legal Entity
|
||||
exercising permissions granted by this License.
|
||||
|
||||
"Source" form shall mean the preferred form for making modifications,
|
||||
including but not limited to software source code, documentation
|
||||
source, and configuration files.
|
||||
|
||||
"Object" form shall mean any form resulting from mechanical
|
||||
transformation or translation of a Source form, including but
|
||||
not limited to compiled object code, generated documentation,
|
||||
and conversions to other media types.
|
||||
|
||||
"Work" shall mean the work of authorship made available under
|
||||
the License, as indicated by a copyright notice that is included in
|
||||
or attached to the work (an example is provided in the Appendix below).
|
||||
|
||||
"Derivative Works" shall mean any work, whether in Source or Object
|
||||
form, that is based on (or derived from) the Work and for which the
|
||||
editorial revisions, annotations, elaborations, or other modifications
|
||||
represent, as a whole, an original work of authorship. For the purposes
|
||||
of this License, Derivative Works shall not include works that remain
|
||||
separable from, or merely link (or bind by name) to the interfaces of,
|
||||
the Work and Derivative Works thereof.
|
||||
|
||||
"Contribution" shall mean, as submitted to the Licensor for inclusion
|
||||
in the Work by the copyright owner or by an individual or Legal Entity
|
||||
authorized to submit on behalf of the copyright owner.
|
||||
|
||||
"Contributor" shall mean Licensor and any Legal Entity on behalf of
|
||||
whom a Contribution has been received by the Licensor and included
|
||||
within the Work.
|
||||
|
||||
2. Grant of Copyright License. Subject to the terms and conditions of
|
||||
this License, each Contributor hereby grants to You a perpetual,
|
||||
worldwide, non-exclusive, no-charge, royalty-free, irrevocable
|
||||
copyright license to reproduce, prepare Derivative Works of,
|
||||
publicly display, publicly perform, sublicense, and distribute the
|
||||
Work and such Derivative Works in Source or Object form.
|
||||
|
||||
3. Grant of Patent License. Subject to the terms and conditions of
|
||||
this License, each Contributor hereby grants to You a perpetual,
|
||||
worldwide, non-exclusive, no-charge, royalty-free, irrevocable
|
||||
(except as stated in this section) patent license to make, have made,
|
||||
use, offer to sell, sell, import, and otherwise transfer the Work,
|
||||
where such license applies only to those patent claims licensable
|
||||
by such Contributor that are necessarily infringed by their
|
||||
Contribution(s) alone or by the combination of their Contribution(s)
|
||||
with the Work to which such Contribution(s) was submitted. If You
|
||||
institute patent litigation against any entity (including a cross-claim
|
||||
or counterclaim in a lawsuit) alleging that the Work or any
|
||||
Contribution embodied within the Work constitutes direct or contributory
|
||||
patent infringement, then any patent licenses granted to You under
|
||||
this License for that Work shall terminate as of the date such
|
||||
litigation is filed.
|
||||
|
||||
4. Redistribution. You may reproduce and distribute copies of the
|
||||
Work or Derivative Works thereof in any medium, with or without
|
||||
modifications, and in Source or Object form, provided that You
|
||||
meet the following conditions:
|
||||
|
||||
(a) You must give any other recipients of the Work or Derivative
|
||||
Works a copy of this License; and
|
||||
|
||||
(b) You must cause any modified files to carry prominent notices
|
||||
stating that You changed the files; and
|
||||
|
||||
(c) You must retain, in the Source form of any Derivative Works
|
||||
that You distribute, all copyright, patent, trademark, and
|
||||
attribution notices from the Source form of the Work; and
|
||||
|
||||
(d) If the Work includes a "NOTICE" text file, You must include a
|
||||
readable copy of the attribution notices contained within such
|
||||
NOTICE file.
|
||||
|
||||
5. Submission of Contributions. Unless You explicitly state otherwise,
|
||||
any Contribution intentionally submitted for inclusion in the Work
|
||||
by You to the Licensor shall be under the terms and conditions of
|
||||
this License, without any additional terms or conditions.
|
||||
|
||||
6. Trademarks. This License does not grant permission to use the trade
|
||||
names, trademarks, service marks, or product names of the Licensor.
|
||||
|
||||
7. Disclaimer of Warranty. Unless required by applicable law or agreed
|
||||
to in writing, Licensor provides the Work (and each Contributor
|
||||
provides its Contributions) on an "AS IS" BASIS, WITHOUT WARRANTIES
|
||||
OR CONDITIONS OF ANY KIND, either express or implied, including,
|
||||
without limitation, any warranties or conditions of TITLE,
|
||||
NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A PARTICULAR
|
||||
PURPOSE. You are solely responsible for determining the
|
||||
appropriateness of using or reproducing the Work and assume any
|
||||
risks associated with Your exercise of permissions under this License.
|
||||
|
||||
8. Limitation of Liability. In no event and under no legal theory,
|
||||
whether in tort (including negligence), contract, or otherwise,
|
||||
unless required by applicable law (such as deliberate and grossly
|
||||
negligent acts) or agreed to in writing, shall any Contributor be
|
||||
liable to You for damages, including any direct, indirect, special,
|
||||
incidental, or exemplary damages of any character arising as a
|
||||
result of this License or out of the use or inability to use the
|
||||
Work (including but not limited to damages for loss of goodwill,
|
||||
work stoppage, computer failure or malfunction, or all other
|
||||
commercial damages or losses), even if such Contributor has been
|
||||
advised of the possibility of such damages.
|
||||
|
||||
9. Accepting Warranty or Additional Liability. While redistributing
|
||||
the Work or Derivative Works thereof, You may choose to offer,
|
||||
and charge a fee for, acceptance of support, warranty, indemnity,
|
||||
or other liability obligations and/or rights consistent with this
|
||||
License. However, in accepting such obligations, You may act only
|
||||
on Your own behalf and on Your sole responsibility, not on behalf
|
||||
of any other Contributor, and only if You agree to indemnify,
|
||||
defend, and hold each Contributor harmless for any liability
|
||||
incurred by, or claims asserted against, such Contributor by reason
|
||||
of your accepting any such warranty or additional liability.
|
||||
|
||||
END OF TERMS AND CONDITIONS
|
||||
|
||||
Copyright 2025 tronghieuit
|
||||
|
||||
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.
|
||||
@@ -0,0 +1,90 @@
|
||||
import os
|
||||
import torch
|
||||
import soundfile as sf
|
||||
from tiny_tts.text.english import normalize_text, grapheme_to_phoneme
|
||||
from tiny_tts.text import phonemes_to_ids
|
||||
from tiny_tts.nn import commons
|
||||
from tiny_tts.models.synthesizer import VoiceSynthesizer
|
||||
from tiny_tts.text.symbols import symbols
|
||||
from tiny_tts.utils.config import (
|
||||
SAMPLING_RATE, SEGMENT_FRAMES, ADD_BLANK, SPEC_CHANNELS,
|
||||
N_SPEAKERS, SPK2ID, MODEL_PARAMS,
|
||||
)
|
||||
from tiny_tts.infer import load_engine
|
||||
|
||||
class TinyTTS:
|
||||
def __init__(self, checkpoint_path=None, device=None):
|
||||
if device is None:
|
||||
self.device = 'cuda' if torch.cuda.is_available() else 'cpu'
|
||||
else:
|
||||
self.device = device
|
||||
|
||||
if checkpoint_path is None:
|
||||
# Look for default checkpoint in pacakage
|
||||
pkg_dir = os.path.dirname(os.path.abspath(__file__))
|
||||
default_ckpt = os.path.join(os.path.dirname(pkg_dir), "checkpoints", "G.pth")
|
||||
# 2. Check HuggingFace Cache / Download
|
||||
if not os.path.exists(default_ckpt):
|
||||
try:
|
||||
from huggingface_hub import hf_hub_download
|
||||
print("Downloading/Loading checkpoint from Hugging Face Hub (backtracking/tiny-tts)...")
|
||||
default_ckpt = hf_hub_download(repo_id="backtracking/tiny-tts", filename="G.pth")
|
||||
except ImportError:
|
||||
raise ImportError("huggingface_hub is required to auto-download the model. Run: pip install huggingface_hub")
|
||||
except Exception as e:
|
||||
raise ValueError(f"Failed to download checkpoint from Hugging Face: {e}")
|
||||
|
||||
checkpoint_path = default_ckpt
|
||||
|
||||
self.model = load_engine(checkpoint_path, self.device)
|
||||
|
||||
def speak(self, text, output_path="output.wav", speaker="MALE", speed=1.0):
|
||||
"""Synthesize text to speech and save to output_path."""
|
||||
print(f"Synthesizing: {text}")
|
||||
|
||||
# Normalize text
|
||||
normalized = normalize_text(text)
|
||||
|
||||
# Phonemize
|
||||
phones, tones, word2ph = grapheme_to_phoneme(normalized)
|
||||
|
||||
# Convert to sequence
|
||||
phone_ids, tone_ids, lang_ids = phonemes_to_ids(phones, tones, "EN")
|
||||
|
||||
# Add blanks
|
||||
if ADD_BLANK:
|
||||
phone_ids = commons.insert_blanks(phone_ids, 0)
|
||||
tone_ids = commons.insert_blanks(tone_ids, 0)
|
||||
lang_ids = commons.insert_blanks(lang_ids, 0)
|
||||
|
||||
x = torch.LongTensor(phone_ids).unsqueeze(0).to(self.device)
|
||||
x_lengths = torch.LongTensor([len(phone_ids)]).to(self.device)
|
||||
tone = torch.LongTensor(tone_ids).unsqueeze(0).to(self.device)
|
||||
language = torch.LongTensor(lang_ids).unsqueeze(0).to(self.device)
|
||||
|
||||
# Speaker ID
|
||||
if speaker not in SPK2ID:
|
||||
print(f"Warning: Speaker '{speaker}' not found, using ID 0. Available: {list(SPK2ID.keys())}")
|
||||
sid = torch.LongTensor([0]).to(self.device)
|
||||
else:
|
||||
sid = torch.LongTensor([SPK2ID[speaker]]).to(self.device)
|
||||
|
||||
# BERT features (disabled - using zero tensors)
|
||||
bert = torch.zeros(1024, len(phone_ids)).to(self.device).unsqueeze(0)
|
||||
ja_bert = torch.zeros(768, len(phone_ids)).to(self.device).unsqueeze(0)
|
||||
|
||||
# speed > 1.0 = faster speech, < 1.0 = slower speech
|
||||
length_scale = 1.0 / speed
|
||||
|
||||
with torch.no_grad():
|
||||
audio, *_ = self.model.infer(
|
||||
x, x_lengths, sid, tone, language, bert, ja_bert,
|
||||
noise_scale=0.667,
|
||||
noise_scale_w=0.8,
|
||||
length_scale=length_scale
|
||||
)
|
||||
|
||||
audio_np = audio[0, 0].cpu().numpy()
|
||||
sf.write(output_path, audio_np, SAMPLING_RATE)
|
||||
print(f"Saved audio to {output_path}")
|
||||
return audio_np
|
||||
@@ -0,0 +1,16 @@
|
||||
from numpy import zeros, int32, float32
|
||||
from torch import from_numpy
|
||||
|
||||
from .core import viterbi_decode_kernel
|
||||
|
||||
|
||||
def viterbi_decode(neg_cent, mask):
|
||||
device = neg_cent.device
|
||||
dtype = neg_cent.dtype
|
||||
neg_cent = neg_cent.data.cpu().numpy().astype(float32)
|
||||
path = zeros(neg_cent.shape, dtype=int32)
|
||||
|
||||
t_t_max = mask.sum(1)[:, 0].data.cpu().numpy().astype(int32)
|
||||
t_s_max = mask.sum(2)[:, 0].data.cpu().numpy().astype(int32)
|
||||
viterbi_decode_kernel(path, neg_cent, t_t_max, t_s_max)
|
||||
return from_numpy(path).to(device=device, dtype=dtype)
|
||||
@@ -0,0 +1,46 @@
|
||||
import numba
|
||||
|
||||
|
||||
@numba.jit(
|
||||
numba.void(
|
||||
numba.int32[:, :, ::1],
|
||||
numba.float32[:, :, ::1],
|
||||
numba.int32[::1],
|
||||
numba.int32[::1],
|
||||
),
|
||||
nopython=True,
|
||||
nogil=True,
|
||||
)
|
||||
def viterbi_decode_kernel(paths, values, t_ys, t_xs):
|
||||
b = paths.shape[0]
|
||||
max_neg_val = -1e9
|
||||
for i in range(int(b)):
|
||||
path = paths[i]
|
||||
value = values[i]
|
||||
t_y = t_ys[i]
|
||||
t_x = t_xs[i]
|
||||
|
||||
v_prev = v_cur = 0.0
|
||||
index = t_x - 1
|
||||
|
||||
for y in range(t_y):
|
||||
for x in range(max(0, t_x + y - t_y), min(t_x, y + 1)):
|
||||
if x == y:
|
||||
v_cur = max_neg_val
|
||||
else:
|
||||
v_cur = value[y - 1, x]
|
||||
if x == 0:
|
||||
if y == 0:
|
||||
v_prev = 0.0
|
||||
else:
|
||||
v_prev = max_neg_val
|
||||
else:
|
||||
v_prev = value[y - 1, x - 1]
|
||||
value[y, x] += max(v_prev, v_cur)
|
||||
|
||||
for y in range(t_y - 1, -1, -1):
|
||||
path[y, index] = 1
|
||||
if index != 0 and (
|
||||
index == y or value[y - 1, index] < value[y - 1, index - 1]
|
||||
):
|
||||
index = index - 1
|
||||
+191
@@ -0,0 +1,191 @@
|
||||
import os
|
||||
import sys
|
||||
import re
|
||||
import torch
|
||||
import soundfile as sf
|
||||
import argparse
|
||||
from tiny_tts.text.english import normalize_text, grapheme_to_phoneme
|
||||
from tiny_tts.text import phonemes_to_ids
|
||||
from tiny_tts.nn import commons
|
||||
from tiny_tts.models import VoiceSynthesizer
|
||||
from tiny_tts.text.symbols import symbols
|
||||
from tiny_tts.utils import (
|
||||
SAMPLING_RATE, SEGMENT_FRAMES, ADD_BLANK, SPEC_CHANNELS,
|
||||
N_SPEAKERS, SPK2ID, MODEL_PARAMS,
|
||||
)
|
||||
|
||||
|
||||
def load_engine(checkpoint_path, device='cuda'):
|
||||
print(f"Loading model from {checkpoint_path}")
|
||||
net_g = VoiceSynthesizer(
|
||||
len(symbols),
|
||||
SPEC_CHANNELS,
|
||||
SEGMENT_FRAMES,
|
||||
n_speakers=N_SPEAKERS,
|
||||
**MODEL_PARAMS
|
||||
).to(device)
|
||||
|
||||
# Count model parameters
|
||||
total_params = sum(p.numel() for p in net_g.parameters())
|
||||
trainable_params = sum(p.numel() for p in net_g.parameters() if p.requires_grad)
|
||||
print(f"Model parameters: {total_params/1e6:.2f}M total, {trainable_params/1e6:.2f}M trainable")
|
||||
|
||||
checkpoint = torch.load(checkpoint_path, map_location=device)
|
||||
state_dict = checkpoint['model']
|
||||
|
||||
# Remove module. prefix and filter shape mismatches
|
||||
model_state = net_g.state_dict()
|
||||
new_state_dict = {}
|
||||
skipped = []
|
||||
for k, v in state_dict.items():
|
||||
key = k[7:] if k.startswith('module.') else k
|
||||
if key in model_state:
|
||||
if v.shape == model_state[key].shape:
|
||||
new_state_dict[key] = v
|
||||
else:
|
||||
skipped.append(f"{key}: ckpt{v.shape} vs model{model_state[key].shape}")
|
||||
else:
|
||||
new_state_dict[key] = v
|
||||
|
||||
if skipped:
|
||||
print(f"Skipped {len(skipped)} mismatched keys:")
|
||||
for s in skipped[:5]:
|
||||
print(f" {s}")
|
||||
if len(skipped) > 5:
|
||||
print(f" ... and {len(skipped)-5} more")
|
||||
|
||||
net_g.load_state_dict(new_state_dict, strict=False)
|
||||
net_g.eval()
|
||||
|
||||
# Fold weight_norm into weight tensors for faster inference (~18% speedup)
|
||||
net_g.dec.remove_weight_norm()
|
||||
|
||||
return net_g
|
||||
|
||||
|
||||
def synthesize(text, output_path, model, speaker="MALE", device='cuda', speed=1.0):
|
||||
print(f"Synthesizing: {text}")
|
||||
|
||||
# Normalize text
|
||||
normalized = normalize_text(text)
|
||||
|
||||
# Phonemize
|
||||
phones, tones, word2ph = grapheme_to_phoneme(normalized)
|
||||
|
||||
# Convert to sequence
|
||||
phone_ids, tone_ids, lang_ids = phonemes_to_ids(phones, tones, "EN")
|
||||
|
||||
# Add blanks
|
||||
if ADD_BLANK:
|
||||
phone_ids = commons.insert_blanks(phone_ids, 0)
|
||||
tone_ids = commons.insert_blanks(tone_ids, 0)
|
||||
lang_ids = commons.insert_blanks(lang_ids, 0)
|
||||
|
||||
x = torch.LongTensor(phone_ids).unsqueeze(0).to(device)
|
||||
x_lengths = torch.LongTensor([len(phone_ids)]).to(device)
|
||||
tone = torch.LongTensor(tone_ids).unsqueeze(0).to(device)
|
||||
language = torch.LongTensor(lang_ids).unsqueeze(0).to(device)
|
||||
|
||||
# Speaker ID
|
||||
if speaker not in SPK2ID:
|
||||
print(f"Warning: Speaker {speaker} not found, using ID 0")
|
||||
sid = torch.LongTensor([0]).to(device)
|
||||
else:
|
||||
sid = torch.LongTensor([SPK2ID[speaker]]).to(device)
|
||||
|
||||
# BERT features (disabled - using zero tensors)
|
||||
bert = torch.zeros(1024, len(phone_ids)).to(device).unsqueeze(0)
|
||||
ja_bert = torch.zeros(768, len(phone_ids)).to(device).unsqueeze(0)
|
||||
|
||||
# speed > 1.0 = faster speech, < 1.0 = slower speech
|
||||
length_scale = 1.0 / speed
|
||||
|
||||
with torch.no_grad():
|
||||
audio, *_ = model.infer(
|
||||
x, x_lengths, sid, tone, language, bert, ja_bert,
|
||||
noise_scale=0.667,
|
||||
noise_scale_w=0.8,
|
||||
length_scale=length_scale
|
||||
)
|
||||
|
||||
audio = audio[0, 0].cpu().numpy()
|
||||
sf.write(output_path, audio, SAMPLING_RATE)
|
||||
print(f"Saved audio to {output_path}")
|
||||
|
||||
|
||||
def get_latest_checkpoint(checkpoint_dir):
|
||||
"""Finds the latest G_*.pth checkpoint in the given directory."""
|
||||
checkpoints = [f for f in os.listdir(checkpoint_dir) if f.startswith('G_') and f.endswith('.pth')]
|
||||
if not checkpoints:
|
||||
return None
|
||||
|
||||
def get_step(filename):
|
||||
match = re.search(r'_(\d+)\.pth', filename)
|
||||
return int(match.group(1)) if match else -1
|
||||
|
||||
latest_ckpt = max(checkpoints, key=get_step)
|
||||
return os.path.join(checkpoint_dir, latest_ckpt)
|
||||
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser(description="TinyTTS — English Text-to-Speech Inference")
|
||||
parser.add_argument("--text", "-t", type=str, default="The weather is nice today, and I feel very relaxed.", help="Text to synthesize")
|
||||
parser.add_argument("--checkpoint", "-c", type=str, default=None, help="Path to checkpoint. Auto-downloads if not provided.")
|
||||
parser.add_argument("--output", "-o", type=str, default="output.wav", help="Output audio file path")
|
||||
parser.add_argument("--speaker", "-s", type=str, default="MALE", help="Speaker ID")
|
||||
parser.add_argument("--speed", type=float, default=1.0, help="Speech speed (1.0=normal, 1.5=faster, 0.7=slower)")
|
||||
parser.add_argument("--device", type=str, default="cuda", help="Device to use (cuda or cpu)")
|
||||
|
||||
args = parser.parse_args()
|
||||
|
||||
if args.checkpoint is None:
|
||||
try:
|
||||
from huggingface_hub import hf_hub_download
|
||||
print("Downloading/Loading checkpoint from Hugging Face Hub (backtracking/tiny-tts)...")
|
||||
args.checkpoint = hf_hub_download(repo_id="backtracking/tiny-tts", filename="G.pth")
|
||||
except ImportError:
|
||||
print("Error: huggingface_hub is required for auto-download. Run: pip install huggingface_hub")
|
||||
sys.exit(1)
|
||||
except Exception as e:
|
||||
print(f"Error downloading checkpoint: {e}")
|
||||
sys.exit(1)
|
||||
|
||||
if not os.path.exists(args.checkpoint):
|
||||
print(f"Error: Checkpoint or directory not found at {args.checkpoint}")
|
||||
sys.exit(1)
|
||||
|
||||
if os.path.isdir(args.checkpoint):
|
||||
latest_ckpt = get_latest_checkpoint(args.checkpoint)
|
||||
if not latest_ckpt:
|
||||
print(f"Error: No G_*.pth checkpoints found in directory {args.checkpoint}")
|
||||
sys.exit(1)
|
||||
args.checkpoint = latest_ckpt
|
||||
print(f"Auto-detected latest checkpoint: {args.checkpoint}")
|
||||
|
||||
# Extract step from checkpoint filename
|
||||
ckpt_basename = os.path.basename(args.checkpoint)
|
||||
match = re.search(r'_(\d+)\.pth', ckpt_basename)
|
||||
step_str = match.group(1) if match else "unknown"
|
||||
|
||||
# Save to output folder
|
||||
out_dir = "infer_outputs"
|
||||
os.makedirs(out_dir, exist_ok=True)
|
||||
|
||||
out_name = os.path.basename(args.output)
|
||||
name, ext = os.path.splitext(out_name)
|
||||
model = load_engine(args.checkpoint, args.device)
|
||||
|
||||
if args.speaker.lower() == "all":
|
||||
if not SPK2ID:
|
||||
print("Error: No speakers found")
|
||||
sys.exit(1)
|
||||
print(f"Synthesizing for all {len(SPK2ID)} speakers...")
|
||||
for spk in SPK2ID.keys():
|
||||
final_output = os.path.join(out_dir, f"{name}_step{step_str}_spk{spk}{ext}")
|
||||
synthesize(args.text, final_output, model, speaker=spk, device=args.device, speed=args.speed)
|
||||
else:
|
||||
final_output = os.path.join(out_dir, f"{name}_step{step_str}_spk{args.speaker}{ext}")
|
||||
synthesize(args.text, final_output, model, speaker=args.speaker, device=args.device, speed=args.speed)
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,199 @@
|
||||
"""
|
||||
ONNX Runtime inference engine for TinyTTS.
|
||||
|
||||
Replaces the PyTorch VoiceSynthesizer.infer() with equivalent
|
||||
ONNX Runtime sessions + NumPy ops for the non-exported parts
|
||||
(alignment path computation).
|
||||
"""
|
||||
import os
|
||||
import numpy as np
|
||||
import soundfile as sf
|
||||
|
||||
from tiny_tts.text.english import normalize_text, grapheme_to_phoneme
|
||||
from tiny_tts.text import phonemes_to_ids
|
||||
from tiny_tts.nn import commons
|
||||
from tiny_tts.utils.config import (
|
||||
SAMPLING_RATE, ADD_BLANK, SPK2ID,
|
||||
)
|
||||
|
||||
try:
|
||||
import onnxruntime as ort
|
||||
except ImportError:
|
||||
raise ImportError("onnxruntime is required. Run: pip install onnxruntime")
|
||||
|
||||
|
||||
def _build_session(path: str, use_gpu: bool = False):
|
||||
"""Create an ORT InferenceSession with optional GPU support."""
|
||||
providers = (
|
||||
["CUDAExecutionProvider", "CPUExecutionProvider"]
|
||||
if use_gpu else
|
||||
["CPUExecutionProvider"]
|
||||
)
|
||||
opts = ort.SessionOptions()
|
||||
opts.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL
|
||||
opts.intra_op_num_threads = os.cpu_count() or 4
|
||||
return ort.InferenceSession(path, sess_options=opts, providers=providers)
|
||||
|
||||
|
||||
def _create_length_mask_np(lengths, max_len=None):
|
||||
"""NumPy equivalent of commons.create_length_mask."""
|
||||
if max_len is None:
|
||||
max_len = int(lengths.max())
|
||||
ids = np.arange(max_len, dtype=np.float32) # [T]
|
||||
mask = (ids[None, :] < lengths[:, None]).astype(np.float32) # [B, T]
|
||||
return mask
|
||||
|
||||
|
||||
def _compute_alignment_path_np(w_ceil, attn_mask):
|
||||
"""
|
||||
Monotonic alignment path - vectorized via cumsum (much faster than Python loops).
|
||||
w_ceil: [B, 1, T_x] — integer duration per phone
|
||||
attn_mask: [B, 1, T_y, T_x] — joint mask
|
||||
Returns attn: [B, 1, T_y, T_x]
|
||||
"""
|
||||
B, _, T_x = w_ceil.shape
|
||||
T_y = attn_mask.shape[2]
|
||||
|
||||
# Build duration matrix: for each phone column expand the duration
|
||||
# cumulative sum of durations gives us the end frame index for each phone
|
||||
dur = w_ceil[:, 0, :] # [B, T_x]
|
||||
cum_dur = np.cumsum(dur, axis=1) # [B, T_x] — end frame (1-indexed)
|
||||
cum_dur_prev = np.pad(cum_dur[:, :-1], ((0,0),(1,0))) # [B, T_x] — start frame
|
||||
|
||||
# Frame indices: [1, T_y, 1]
|
||||
frame_idx = np.arange(T_y, dtype=np.float32)[None, :, None] # [1, T_y, 1]
|
||||
# For each phone, mark frames [start, end)
|
||||
# cum_dur_prev: [B,1,T_x], cum_dur: [B,1,T_x]
|
||||
start = cum_dur_prev[:, None, :] # [B, 1, T_x]
|
||||
end = cum_dur[:, None, :] # [B, 1, T_x]
|
||||
attn = ((frame_idx >= start) & (frame_idx < end)).astype(np.float32) # [B, T_y, T_x]
|
||||
attn = attn[:, None, :, :] # [B, 1, T_y, T_x]
|
||||
return attn * attn_mask
|
||||
|
||||
|
||||
class OnnxTinyTTS:
|
||||
"""
|
||||
Inference using ONNX Runtime.
|
||||
|
||||
Args:
|
||||
onnx_dir: directory containing the 4 .onnx files
|
||||
use_gpu: if True, try CUDAExecutionProvider
|
||||
"""
|
||||
|
||||
def __init__(self, onnx_dir: str = "onnx", use_gpu: bool = False):
|
||||
onnx_dir = os.path.abspath(onnx_dir)
|
||||
print(f"Loading ONNX sessions from: {onnx_dir}")
|
||||
|
||||
self._enc = _build_session(os.path.join(onnx_dir, "text_encoder.onnx"), use_gpu)
|
||||
self._dp = _build_session(os.path.join(onnx_dir, "duration_predictor.onnx"), use_gpu)
|
||||
self._flow = _build_session(os.path.join(onnx_dir, "flow.onnx"), use_gpu)
|
||||
self._dec = _build_session(os.path.join(onnx_dir, "decoder.onnx"), use_gpu)
|
||||
|
||||
print("ONNX sessions ready ✅")
|
||||
|
||||
def _text_to_ids(self, text: str):
|
||||
normalized = normalize_text(text)
|
||||
phones, tones, _ = grapheme_to_phoneme(normalized)
|
||||
phone_ids, tone_ids, lang_ids = phonemes_to_ids(phones, tones, "EN")
|
||||
|
||||
if ADD_BLANK:
|
||||
phone_ids = commons.insert_blanks(phone_ids, 0)
|
||||
tone_ids = commons.insert_blanks(tone_ids, 0)
|
||||
lang_ids = commons.insert_blanks(lang_ids, 0)
|
||||
|
||||
return phone_ids, tone_ids, lang_ids
|
||||
|
||||
def speak(
|
||||
self,
|
||||
text: str,
|
||||
output_path: str = "onnx_output.wav",
|
||||
speaker: str = "female",
|
||||
noise_scale: float = 0.667,
|
||||
noise_scale_w: float = 0.8,
|
||||
length_scale: float = 1.0,
|
||||
output_sr: int = None,
|
||||
) -> np.ndarray:
|
||||
"""Synthesize speech and save to output_path.
|
||||
|
||||
Args:
|
||||
output_sr: If set (e.g. 22050), resample the output from 44100 Hz.
|
||||
Useful to reduce file size while keeping quality.
|
||||
"""
|
||||
print(f"[ONNX] Synthesizing: {text}")
|
||||
|
||||
phone_ids, tone_ids, lang_ids = self._text_to_ids(text)
|
||||
T = len(phone_ids)
|
||||
|
||||
# Prepare inputs as float32 / int64 arrays
|
||||
x = np.array(phone_ids, dtype=np.int64)[None, :] # [1, T]
|
||||
x_len = np.array([T], dtype=np.int64) # [1]
|
||||
tone = np.array(tone_ids, dtype=np.int64)[None, :] # [1, T]
|
||||
lang = np.array(lang_ids, dtype=np.int64)[None, :] # [1, T]
|
||||
bert = np.zeros((1, 1024, T), dtype=np.float32)
|
||||
ja_bert = np.zeros((1, 768, T), dtype=np.float32)
|
||||
sid_val = SPK2ID.get(speaker, 0)
|
||||
sid = np.array([sid_val], dtype=np.int64) # [1]
|
||||
|
||||
# ── 1. Text Encoder ──────────────────────────────────────────────
|
||||
x_enc, m_p, logs_p, x_mask, g = self._enc.run(
|
||||
None,
|
||||
{
|
||||
"phone_ids": x,
|
||||
"phone_lengths":x_len,
|
||||
"tone_ids": tone,
|
||||
"language_ids": lang,
|
||||
"bert": bert,
|
||||
"ja_bert": ja_bert,
|
||||
"speaker_id": sid,
|
||||
},
|
||||
)
|
||||
|
||||
# ── 2. Duration Predictor ─────────────────────────────────────────
|
||||
logw = self._dp.run(None, {"x": x_enc, "x_mask": x_mask, "g": g})[0]
|
||||
|
||||
# ── 3. Alignment Path (NumPy) ─────────────────────────────────────
|
||||
w = np.exp(logw) * x_mask * length_scale # [1, 1, T]
|
||||
w_ceil = np.ceil(w) # [1, 1, T]
|
||||
y_len = max(1, int(w_ceil.sum()))
|
||||
y_lens = np.array([y_len], dtype=np.int64)
|
||||
|
||||
y_mask = _create_length_mask_np(y_lens, y_len) # [1, T_y]
|
||||
y_mask = y_mask[:, None, :] # [1, 1, T_y]
|
||||
# attn_mask: [1, 1, T_y, T_x] (outer product of frame mask and phone mask)
|
||||
attn_mask = y_mask[:, :, :, None] * x_mask[:, :, None, :] # [1,1,T_y,T_x]
|
||||
attn = _compute_alignment_path_np(w_ceil, attn_mask) # [1, 1, T_y, T_x]
|
||||
|
||||
# Expand prior stats via alignment
|
||||
m_p_exp = np.matmul(attn[:, 0], m_p.transpose(0, 2, 1)).transpose(0, 2, 1)
|
||||
logs_p_exp = np.matmul(attn[:, 0], logs_p.transpose(0, 2, 1)).transpose(0, 2, 1)
|
||||
|
||||
# ── 4. Sample z_p ─────────────────────────────────────────────────
|
||||
z_p = m_p_exp + np.random.randn(*m_p_exp.shape).astype(np.float32) * \
|
||||
np.exp(logs_p_exp) * noise_scale
|
||||
|
||||
# ── 5. Flow (reverse) ─────────────────────────────────────────────
|
||||
z = self._flow.run(
|
||||
None,
|
||||
{"z_p": z_p, "y_mask": y_mask.astype(np.float32), "g": g},
|
||||
)[0]
|
||||
|
||||
# ── 6. Decoder ────────────────────────────────────────────────────
|
||||
z_masked = (z * y_mask).astype(np.float32)
|
||||
audio = self._dec.run(None, {"z": z_masked, "g": g})[0] # [1, 1, samples]
|
||||
|
||||
audio_np = audio[0, 0]
|
||||
save_sr = SAMPLING_RATE
|
||||
if output_sr is not None and output_sr != SAMPLING_RATE:
|
||||
try:
|
||||
import torchaudio
|
||||
import torch
|
||||
wav_t = torch.from_numpy(audio_np).unsqueeze(0)
|
||||
resampler = torchaudio.transforms.Resample(SAMPLING_RATE, output_sr)
|
||||
audio_np = resampler(wav_t).squeeze(0).numpy()
|
||||
save_sr = output_sr
|
||||
except Exception as e:
|
||||
print(f"[ONNX] Resampling failed ({e}), saving at {SAMPLING_RATE}Hz")
|
||||
|
||||
sf.write(output_path, audio_np, save_sr)
|
||||
print(f"[ONNX] Saved: {output_path} ({save_sr}Hz)")
|
||||
return audio_np
|
||||
@@ -0,0 +1 @@
|
||||
from .synthesizer import VoiceSynthesizer
|
||||
@@ -0,0 +1,718 @@
|
||||
import math
|
||||
import torch
|
||||
from torch import nn
|
||||
from torch.nn import functional as F
|
||||
|
||||
from tiny_tts.nn import commons
|
||||
from tiny_tts.nn import modules
|
||||
from tiny_tts.nn import attentions
|
||||
|
||||
from torch.nn import Conv1d, ConvTranspose1d
|
||||
from torch.nn.utils import weight_norm, remove_weight_norm
|
||||
|
||||
from tiny_tts.nn.commons import initialize_weights, compute_padding
|
||||
import tiny_tts.alignment as alignment
|
||||
|
||||
|
||||
class AttentionFlowBlock(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
channels,
|
||||
hidden_channels,
|
||||
filter_channels,
|
||||
n_heads,
|
||||
n_layers,
|
||||
kernel_size,
|
||||
p_dropout,
|
||||
n_flows=4,
|
||||
gin_channels=0,
|
||||
share_parameter=False,
|
||||
):
|
||||
super().__init__()
|
||||
self.channels = channels
|
||||
self.hidden_channels = hidden_channels
|
||||
self.kernel_size = kernel_size
|
||||
self.n_layers = n_layers
|
||||
self.n_flows = n_flows
|
||||
self.gin_channels = gin_channels
|
||||
|
||||
self.flows = nn.ModuleList()
|
||||
|
||||
self.wn = (
|
||||
attentions.FeedForward(
|
||||
hidden_channels,
|
||||
filter_channels,
|
||||
n_heads,
|
||||
n_layers,
|
||||
kernel_size,
|
||||
p_dropout,
|
||||
isflow=True,
|
||||
gin_channels=self.gin_channels,
|
||||
)
|
||||
if share_parameter
|
||||
else None
|
||||
)
|
||||
|
||||
for i in range(n_flows):
|
||||
self.flows.append(
|
||||
modules.TransformerCouplingLayer(
|
||||
channels,
|
||||
hidden_channels,
|
||||
kernel_size,
|
||||
n_layers,
|
||||
n_heads,
|
||||
p_dropout,
|
||||
filter_channels,
|
||||
mean_only=True,
|
||||
wn_sharing_parameter=self.wn,
|
||||
gin_channels=self.gin_channels,
|
||||
)
|
||||
)
|
||||
self.flows.append(modules.FlipTransform())
|
||||
|
||||
def forward(self, x, x_mask, g=None, reverse=False):
|
||||
if not reverse:
|
||||
for flow in self.flows:
|
||||
x, _ = flow(x, x_mask, g=g, reverse=reverse)
|
||||
else:
|
||||
for flow in reversed(self.flows):
|
||||
x = flow(x, x_mask, g=g, reverse=reverse)
|
||||
return x
|
||||
|
||||
|
||||
class VariationalDurationModel(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
in_channels,
|
||||
filter_channels,
|
||||
kernel_size,
|
||||
p_dropout,
|
||||
n_flows=4,
|
||||
gin_channels=0,
|
||||
):
|
||||
super().__init__()
|
||||
filter_channels = in_channels
|
||||
self.in_channels = in_channels
|
||||
self.filter_channels = filter_channels
|
||||
self.kernel_size = kernel_size
|
||||
self.p_dropout = p_dropout
|
||||
self.n_flows = n_flows
|
||||
self.gin_channels = gin_channels
|
||||
|
||||
self.log_flow = modules.LogTransform()
|
||||
self.flows = nn.ModuleList()
|
||||
self.flows.append(modules.AffineCoupling(2))
|
||||
for i in range(n_flows):
|
||||
self.flows.append(
|
||||
modules.ConvolutionalFlow(2, filter_channels, kernel_size, n_layers=3)
|
||||
)
|
||||
self.flows.append(modules.FlipTransform())
|
||||
|
||||
self.post_pre = nn.Conv1d(1, filter_channels, 1)
|
||||
self.post_proj = nn.Conv1d(filter_channels, filter_channels, 1)
|
||||
self.post_convs = modules.DepthwiseSepConv(
|
||||
filter_channels, kernel_size, n_layers=3, p_dropout=p_dropout
|
||||
)
|
||||
self.post_flows = nn.ModuleList()
|
||||
self.post_flows.append(modules.AffineCoupling(2))
|
||||
for i in range(4):
|
||||
self.post_flows.append(
|
||||
modules.ConvolutionalFlow(2, filter_channels, kernel_size, n_layers=3)
|
||||
)
|
||||
self.post_flows.append(modules.FlipTransform())
|
||||
|
||||
self.pre = nn.Conv1d(in_channels, filter_channels, 1)
|
||||
self.proj = nn.Conv1d(filter_channels, filter_channels, 1)
|
||||
self.convs = modules.DepthwiseSepConv(
|
||||
filter_channels, kernel_size, n_layers=3, p_dropout=p_dropout
|
||||
)
|
||||
if gin_channels != 0:
|
||||
self.cond = nn.Conv1d(gin_channels, filter_channels, 1)
|
||||
|
||||
def forward(self, x, x_mask, w=None, g=None, reverse=False, noise_scale=1.0):
|
||||
x = torch.detach(x)
|
||||
x = self.pre(x)
|
||||
if g is not None:
|
||||
g = torch.detach(g)
|
||||
x = x + self.cond(g)
|
||||
x = self.convs(x, x_mask)
|
||||
x = self.proj(x) * x_mask
|
||||
|
||||
if not reverse:
|
||||
flows = self.flows
|
||||
assert w is not None
|
||||
|
||||
logdet_tot_q = 0
|
||||
h_w = self.post_pre(w)
|
||||
h_w = self.post_convs(h_w, x_mask)
|
||||
h_w = self.post_proj(h_w) * x_mask
|
||||
e_q = (
|
||||
torch.randn(w.size(0), 2, w.size(2)).to(device=x.device, dtype=x.dtype)
|
||||
* x_mask
|
||||
)
|
||||
z_q = e_q
|
||||
for flow in self.post_flows:
|
||||
z_q, logdet_q = flow(z_q, x_mask, g=(x + h_w))
|
||||
logdet_tot_q += logdet_q
|
||||
z_u, z1 = torch.split(z_q, [1, 1], 1)
|
||||
u = torch.sigmoid(z_u) * x_mask
|
||||
z0 = (w - u) * x_mask
|
||||
logdet_tot_q += torch.sum(
|
||||
(F.logsigmoid(z_u) + F.logsigmoid(-z_u)) * x_mask, [1, 2]
|
||||
)
|
||||
logq = (
|
||||
torch.sum(-0.5 * (math.log(2 * math.pi) + (e_q**2)) * x_mask, [1, 2])
|
||||
- logdet_tot_q
|
||||
)
|
||||
|
||||
logdet_tot = 0
|
||||
z0, logdet = self.log_flow(z0, x_mask)
|
||||
logdet_tot += logdet
|
||||
z = torch.cat([z0, z1], 1)
|
||||
for flow in flows:
|
||||
z, logdet = flow(z, x_mask, g=x, reverse=reverse)
|
||||
logdet_tot = logdet_tot + logdet
|
||||
nll = (
|
||||
torch.sum(0.5 * (math.log(2 * math.pi) + (z**2)) * x_mask, [1, 2])
|
||||
- logdet_tot
|
||||
)
|
||||
return nll + logq
|
||||
else:
|
||||
flows = list(reversed(self.flows))
|
||||
flows = flows[:-2] + [flows[-1]]
|
||||
z = (
|
||||
torch.randn(x.size(0), 2, x.size(2)).to(device=x.device, dtype=x.dtype)
|
||||
* noise_scale
|
||||
)
|
||||
for flow in flows:
|
||||
z = flow(z, x_mask, g=x, reverse=reverse)
|
||||
z0, z1 = torch.split(z, [1, 1], 1)
|
||||
logw = z0
|
||||
return logw
|
||||
|
||||
|
||||
class DurationEstimator(nn.Module):
|
||||
def __init__(
|
||||
self, in_channels, filter_channels, kernel_size, p_dropout, gin_channels=0
|
||||
):
|
||||
super().__init__()
|
||||
|
||||
self.in_channels = in_channels
|
||||
self.filter_channels = filter_channels
|
||||
self.kernel_size = kernel_size
|
||||
self.p_dropout = p_dropout
|
||||
self.gin_channels = gin_channels
|
||||
|
||||
self.drop = nn.Dropout(p_dropout)
|
||||
self.conv_1 = nn.Conv1d(
|
||||
in_channels, filter_channels, kernel_size, padding=kernel_size // 2
|
||||
)
|
||||
self.norm_1 = modules.ChannelNorm(filter_channels)
|
||||
self.conv_2 = nn.Conv1d(
|
||||
filter_channels, filter_channels, kernel_size, padding=kernel_size // 2
|
||||
)
|
||||
self.norm_2 = modules.ChannelNorm(filter_channels)
|
||||
self.proj = nn.Conv1d(filter_channels, 1, 1)
|
||||
|
||||
if gin_channels != 0:
|
||||
self.cond = nn.Conv1d(gin_channels, in_channels, 1)
|
||||
|
||||
def forward(self, x, x_mask, g=None):
|
||||
x = torch.detach(x)
|
||||
if g is not None:
|
||||
g = torch.detach(g)
|
||||
x = x + self.cond(g)
|
||||
x = self.conv_1(x * x_mask)
|
||||
x = torch.relu(x)
|
||||
x = self.norm_1(x)
|
||||
x = self.drop(x)
|
||||
x = self.conv_2(x * x_mask)
|
||||
x = torch.relu(x)
|
||||
x = self.norm_2(x)
|
||||
x = self.drop(x)
|
||||
x = self.proj(x * x_mask)
|
||||
return x * x_mask
|
||||
|
||||
|
||||
class PhonemeEncoder(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
n_vocab,
|
||||
out_channels,
|
||||
hidden_channels,
|
||||
filter_channels,
|
||||
n_heads,
|
||||
n_layers,
|
||||
kernel_size,
|
||||
p_dropout,
|
||||
gin_channels=0,
|
||||
num_languages=None,
|
||||
num_tones=None,
|
||||
):
|
||||
super().__init__()
|
||||
if num_languages is None:
|
||||
from tiny_tts.text import num_languages
|
||||
if num_tones is None:
|
||||
from tiny_tts.text import num_tones
|
||||
self.n_vocab = n_vocab
|
||||
self.out_channels = out_channels
|
||||
self.hidden_channels = hidden_channels
|
||||
self.filter_channels = filter_channels
|
||||
self.n_heads = n_heads
|
||||
self.n_layers = n_layers
|
||||
self.kernel_size = kernel_size
|
||||
self.p_dropout = p_dropout
|
||||
self.gin_channels = gin_channels
|
||||
self.emb = nn.Embedding(n_vocab, hidden_channels)
|
||||
nn.init.normal_(self.emb.weight, 0.0, hidden_channels**-0.5)
|
||||
self.tone_emb = nn.Embedding(num_tones, hidden_channels)
|
||||
nn.init.normal_(self.tone_emb.weight, 0.0, hidden_channels**-0.5)
|
||||
self.language_emb = nn.Embedding(num_languages, hidden_channels)
|
||||
nn.init.normal_(self.language_emb.weight, 0.0, hidden_channels**-0.5)
|
||||
self.bert_proj = nn.Conv1d(1024, hidden_channels, 1)
|
||||
self.ja_bert_proj = nn.Conv1d(768, hidden_channels, 1)
|
||||
|
||||
self.encoder = attentions.TransformerBlock(
|
||||
hidden_channels,
|
||||
filter_channels,
|
||||
n_heads,
|
||||
n_layers,
|
||||
kernel_size,
|
||||
p_dropout,
|
||||
gin_channels=self.gin_channels,
|
||||
)
|
||||
self.proj = nn.Conv1d(hidden_channels, out_channels * 2, 1)
|
||||
|
||||
def forward(self, x, x_lengths, tone, language, bert, ja_bert, g=None):
|
||||
bert_emb = self.bert_proj(bert).transpose(1, 2)
|
||||
ja_bert_emb = self.ja_bert_proj(ja_bert).transpose(1, 2)
|
||||
x = (
|
||||
self.emb(x)
|
||||
+ self.tone_emb(tone)
|
||||
+ self.language_emb(language)
|
||||
+ bert_emb
|
||||
+ ja_bert_emb
|
||||
) * math.sqrt(
|
||||
self.hidden_channels
|
||||
)
|
||||
x = torch.transpose(x, 1, -1)
|
||||
x_mask = torch.unsqueeze(commons.create_length_mask(x_lengths, x.size(2)), 1).to(
|
||||
x.dtype
|
||||
)
|
||||
|
||||
x = self.encoder(x * x_mask, x_mask, g=g)
|
||||
stats = self.proj(x) * x_mask
|
||||
|
||||
m, logs = torch.split(stats, self.out_channels, dim=1)
|
||||
return x, m, logs, x_mask
|
||||
|
||||
|
||||
class FlowBlock(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
channels,
|
||||
hidden_channels,
|
||||
kernel_size,
|
||||
dilation_rate,
|
||||
n_layers,
|
||||
n_flows=4,
|
||||
gin_channels=0,
|
||||
):
|
||||
super().__init__()
|
||||
self.channels = channels
|
||||
self.hidden_channels = hidden_channels
|
||||
self.kernel_size = kernel_size
|
||||
self.dilation_rate = dilation_rate
|
||||
self.n_layers = n_layers
|
||||
self.n_flows = n_flows
|
||||
self.gin_channels = gin_channels
|
||||
|
||||
self.flows = nn.ModuleList()
|
||||
for i in range(n_flows):
|
||||
self.flows.append(
|
||||
modules.FlowCouplingLayer(
|
||||
channels,
|
||||
hidden_channels,
|
||||
kernel_size,
|
||||
dilation_rate,
|
||||
n_layers,
|
||||
gin_channels=gin_channels,
|
||||
mean_only=True,
|
||||
)
|
||||
)
|
||||
self.flows.append(modules.FlipTransform())
|
||||
|
||||
def forward(self, x, x_mask, g=None, reverse=False):
|
||||
if not reverse:
|
||||
for flow in self.flows:
|
||||
x, _ = flow(x, x_mask, g=g, reverse=reverse)
|
||||
else:
|
||||
for flow in reversed(self.flows):
|
||||
x = flow(x, x_mask, g=g, reverse=reverse)
|
||||
return x
|
||||
|
||||
|
||||
class LatentEncoder(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
in_channels,
|
||||
out_channels,
|
||||
hidden_channels,
|
||||
kernel_size,
|
||||
dilation_rate,
|
||||
n_layers,
|
||||
gin_channels=0,
|
||||
):
|
||||
super().__init__()
|
||||
self.in_channels = in_channels
|
||||
self.out_channels = out_channels
|
||||
self.hidden_channels = hidden_channels
|
||||
self.kernel_size = kernel_size
|
||||
self.dilation_rate = dilation_rate
|
||||
self.n_layers = n_layers
|
||||
self.gin_channels = gin_channels
|
||||
|
||||
self.pre = nn.Conv1d(in_channels, hidden_channels, 1)
|
||||
self.enc = modules.WaveNet(
|
||||
hidden_channels,
|
||||
kernel_size,
|
||||
dilation_rate,
|
||||
n_layers,
|
||||
gin_channels=gin_channels,
|
||||
)
|
||||
self.proj = nn.Conv1d(hidden_channels, out_channels * 2, 1)
|
||||
|
||||
def forward(self, x, x_lengths, g=None, tau=1.0):
|
||||
x_mask = torch.unsqueeze(commons.create_length_mask(x_lengths, x.size(2)), 1).to(
|
||||
x.dtype
|
||||
)
|
||||
x = self.pre(x) * x_mask
|
||||
x = self.enc(x, x_mask, g=g)
|
||||
stats = self.proj(x) * x_mask
|
||||
m, logs = torch.split(stats, self.out_channels, dim=1)
|
||||
z = (m + torch.randn_like(m) * tau * torch.exp(logs)) * x_mask
|
||||
return z, m, logs, x_mask
|
||||
|
||||
|
||||
class WaveformDecoder(torch.nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
initial_channel,
|
||||
resblock,
|
||||
resblock_kernel_sizes,
|
||||
resblock_dilation_sizes,
|
||||
upsample_rates,
|
||||
upsample_initial_channel,
|
||||
upsample_kernel_sizes,
|
||||
gin_channels=0,
|
||||
):
|
||||
super(WaveformDecoder, self).__init__()
|
||||
self.num_kernels = len(resblock_kernel_sizes)
|
||||
self.num_upsamples = len(upsample_rates)
|
||||
self.conv_pre = Conv1d(
|
||||
initial_channel, upsample_initial_channel, 7, 1, padding=3
|
||||
)
|
||||
resblock = modules.ConvResBlock if resblock == "1" else modules.ConvResBlockLight
|
||||
|
||||
self.ups = nn.ModuleList()
|
||||
for i, (u, k) in enumerate(zip(upsample_rates, upsample_kernel_sizes)):
|
||||
self.ups.append(
|
||||
weight_norm(
|
||||
ConvTranspose1d(
|
||||
upsample_initial_channel // (2**i),
|
||||
upsample_initial_channel // (2 ** (i + 1)),
|
||||
k,
|
||||
u,
|
||||
padding=(k - u) // 2,
|
||||
)
|
||||
)
|
||||
)
|
||||
|
||||
self.resblocks = nn.ModuleList()
|
||||
for i in range(len(self.ups)):
|
||||
ch = upsample_initial_channel // (2 ** (i + 1))
|
||||
for j, (k, d) in enumerate(
|
||||
zip(resblock_kernel_sizes, resblock_dilation_sizes)
|
||||
):
|
||||
self.resblocks.append(resblock(ch, k, d))
|
||||
|
||||
self.conv_post = Conv1d(ch, 1, 7, 1, padding=3, bias=False)
|
||||
self.ups.apply(initialize_weights)
|
||||
|
||||
if gin_channels != 0:
|
||||
self.cond = nn.Conv1d(gin_channels, upsample_initial_channel, 1)
|
||||
|
||||
def forward(self, x, g=None):
|
||||
x = self.conv_pre(x)
|
||||
if g is not None:
|
||||
x = x + self.cond(g)
|
||||
|
||||
for i in range(self.num_upsamples):
|
||||
x = F.leaky_relu(x, modules.LRELU_SLOPE)
|
||||
x = self.ups[i](x)
|
||||
xs = None
|
||||
for j in range(self.num_kernels):
|
||||
if xs is None:
|
||||
xs = self.resblocks[i * self.num_kernels + j](x)
|
||||
else:
|
||||
xs += self.resblocks[i * self.num_kernels + j](x)
|
||||
x = xs / self.num_kernels
|
||||
x = F.leaky_relu(x)
|
||||
x = self.conv_post(x)
|
||||
x = torch.tanh(x)
|
||||
|
||||
return x
|
||||
|
||||
def remove_weight_norm(self):
|
||||
for layer in self.ups:
|
||||
remove_weight_norm(layer)
|
||||
for layer in self.resblocks:
|
||||
layer.remove_weight_norm()
|
||||
|
||||
|
||||
class StyleEncoder(nn.Module):
|
||||
def __init__(self, spec_channels, gin_channels=0, layernorm=False):
|
||||
super().__init__()
|
||||
self.spec_channels = spec_channels
|
||||
ref_enc_filters = [32, 32, 64, 64, 128, 128]
|
||||
K = len(ref_enc_filters)
|
||||
filters = [1] + ref_enc_filters
|
||||
convs = [
|
||||
weight_norm(
|
||||
nn.Conv2d(
|
||||
in_channels=filters[i],
|
||||
out_channels=filters[i + 1],
|
||||
kernel_size=(3, 3),
|
||||
stride=(2, 2),
|
||||
padding=(1, 1),
|
||||
)
|
||||
)
|
||||
for i in range(K)
|
||||
]
|
||||
self.convs = nn.ModuleList(convs)
|
||||
|
||||
out_channels = self.calculate_channels(spec_channels, 3, 2, 1, K)
|
||||
self.gru = nn.GRU(
|
||||
input_size=ref_enc_filters[-1] * out_channels,
|
||||
hidden_size=256 // 2,
|
||||
batch_first=True,
|
||||
)
|
||||
self.proj = nn.Linear(128, gin_channels)
|
||||
if layernorm:
|
||||
self.layernorm = nn.LayerNorm(self.spec_channels)
|
||||
else:
|
||||
self.layernorm = None
|
||||
|
||||
def forward(self, inputs, mask=None):
|
||||
N = inputs.size(0)
|
||||
|
||||
out = inputs.view(N, 1, -1, self.spec_channels)
|
||||
if self.layernorm is not None:
|
||||
out = self.layernorm(out)
|
||||
|
||||
for conv in self.convs:
|
||||
out = conv(out)
|
||||
out = F.relu(out)
|
||||
|
||||
out = out.transpose(1, 2)
|
||||
T = out.size(1)
|
||||
N = out.size(0)
|
||||
out = out.contiguous().view(N, T, -1)
|
||||
|
||||
self.gru.flatten_parameters()
|
||||
memory, out = self.gru(out)
|
||||
|
||||
return self.proj(out.squeeze(0))
|
||||
|
||||
def calculate_channels(self, L, kernel_size, stride, pad, n_convs):
|
||||
for i in range(n_convs):
|
||||
L = (L - kernel_size + 2 * pad) // stride + 1
|
||||
return L
|
||||
|
||||
|
||||
class VoiceSynthesizer(nn.Module):
|
||||
"""Voice synthesis model for inference."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
n_vocab,
|
||||
spec_channels,
|
||||
segment_size,
|
||||
inter_channels,
|
||||
hidden_channels,
|
||||
filter_channels,
|
||||
n_heads,
|
||||
n_layers,
|
||||
kernel_size,
|
||||
p_dropout,
|
||||
resblock,
|
||||
resblock_kernel_sizes,
|
||||
resblock_dilation_sizes,
|
||||
upsample_rates,
|
||||
upsample_initial_channel,
|
||||
upsample_kernel_sizes,
|
||||
n_speakers=256,
|
||||
gin_channels=256,
|
||||
use_sdp=True,
|
||||
n_flow_layer=4,
|
||||
n_layers_trans_flow=6,
|
||||
flow_share_parameter=False,
|
||||
use_transformer_flow=True,
|
||||
use_vc=False,
|
||||
num_languages=None,
|
||||
num_tones=None,
|
||||
norm_refenc=False,
|
||||
**kwargs
|
||||
):
|
||||
super().__init__()
|
||||
self.n_vocab = n_vocab
|
||||
self.spec_channels = spec_channels
|
||||
self.inter_channels = inter_channels
|
||||
self.hidden_channels = hidden_channels
|
||||
self.filter_channels = filter_channels
|
||||
self.n_heads = n_heads
|
||||
self.n_layers = n_layers
|
||||
self.kernel_size = kernel_size
|
||||
self.p_dropout = p_dropout
|
||||
self.resblock = resblock
|
||||
self.resblock_kernel_sizes = resblock_kernel_sizes
|
||||
self.resblock_dilation_sizes = resblock_dilation_sizes
|
||||
self.upsample_rates = upsample_rates
|
||||
self.upsample_initial_channel = upsample_initial_channel
|
||||
self.upsample_kernel_sizes = upsample_kernel_sizes
|
||||
self.segment_size = segment_size
|
||||
self.n_speakers = n_speakers
|
||||
self.gin_channels = gin_channels
|
||||
self.n_layers_trans_flow = n_layers_trans_flow
|
||||
self.use_spk_conditioned_encoder = kwargs.get(
|
||||
"use_spk_conditioned_encoder", True
|
||||
)
|
||||
self.use_sdp = use_sdp
|
||||
self.use_noise_scaled_mas = kwargs.get("use_noise_scaled_mas", False)
|
||||
self.mas_noise_scale_initial = kwargs.get("mas_noise_scale_initial", 0.01)
|
||||
self.noise_scale_delta = kwargs.get("noise_scale_delta", 2e-6)
|
||||
self.current_mas_noise_scale = self.mas_noise_scale_initial
|
||||
if self.use_spk_conditioned_encoder and gin_channels > 0:
|
||||
self.enc_gin_channels = gin_channels
|
||||
else:
|
||||
self.enc_gin_channels = 0
|
||||
self.enc_p = PhonemeEncoder(
|
||||
n_vocab,
|
||||
inter_channels,
|
||||
hidden_channels,
|
||||
filter_channels,
|
||||
n_heads,
|
||||
n_layers,
|
||||
kernel_size,
|
||||
p_dropout,
|
||||
gin_channels=self.enc_gin_channels,
|
||||
num_languages=num_languages,
|
||||
num_tones=num_tones,
|
||||
)
|
||||
self.dec = WaveformDecoder(
|
||||
inter_channels,
|
||||
resblock,
|
||||
resblock_kernel_sizes,
|
||||
resblock_dilation_sizes,
|
||||
upsample_rates,
|
||||
upsample_initial_channel,
|
||||
upsample_kernel_sizes,
|
||||
gin_channels=gin_channels,
|
||||
)
|
||||
self.enc_q = LatentEncoder(
|
||||
spec_channels,
|
||||
inter_channels,
|
||||
hidden_channels,
|
||||
5,
|
||||
1,
|
||||
16,
|
||||
gin_channels=gin_channels,
|
||||
)
|
||||
if use_transformer_flow:
|
||||
self.flow = AttentionFlowBlock(
|
||||
inter_channels,
|
||||
hidden_channels,
|
||||
filter_channels,
|
||||
n_heads,
|
||||
n_layers_trans_flow,
|
||||
5,
|
||||
p_dropout,
|
||||
n_flow_layer,
|
||||
gin_channels=gin_channels,
|
||||
share_parameter=flow_share_parameter,
|
||||
)
|
||||
else:
|
||||
self.flow = FlowBlock(
|
||||
inter_channels,
|
||||
hidden_channels,
|
||||
5,
|
||||
1,
|
||||
n_flow_layer,
|
||||
gin_channels=gin_channels,
|
||||
)
|
||||
self.sdp = VariationalDurationModel(
|
||||
hidden_channels, 192, 3, 0.5, 4, gin_channels=gin_channels
|
||||
)
|
||||
self.dp = DurationEstimator(
|
||||
hidden_channels, 256, 3, 0.5, gin_channels=gin_channels
|
||||
)
|
||||
|
||||
if n_speakers > 0:
|
||||
self.emb_g = nn.Embedding(n_speakers, gin_channels)
|
||||
else:
|
||||
self.ref_enc = StyleEncoder(spec_channels, gin_channels, layernorm=norm_refenc)
|
||||
self.use_vc = use_vc
|
||||
|
||||
def infer(
|
||||
self,
|
||||
x,
|
||||
x_lengths,
|
||||
sid,
|
||||
tone,
|
||||
language,
|
||||
bert,
|
||||
ja_bert,
|
||||
noise_scale=0.667,
|
||||
length_scale=1,
|
||||
noise_scale_w=0.8,
|
||||
max_len=None,
|
||||
sdp_ratio=0,
|
||||
y=None,
|
||||
g=None,
|
||||
):
|
||||
if g is None:
|
||||
if self.n_speakers > 0:
|
||||
g = self.emb_g(sid).unsqueeze(-1)
|
||||
else:
|
||||
g = self.ref_enc(y.transpose(1, 2)).unsqueeze(-1)
|
||||
if self.use_vc:
|
||||
g_p = None
|
||||
else:
|
||||
g_p = g
|
||||
x, m_p, logs_p, x_mask = self.enc_p(
|
||||
x, x_lengths, tone, language, bert, ja_bert, g=g_p
|
||||
)
|
||||
logw = self.sdp(x, x_mask, g=g, reverse=True, noise_scale=noise_scale_w) * (
|
||||
sdp_ratio
|
||||
) + self.dp(x, x_mask, g=g) * (1 - sdp_ratio)
|
||||
w = torch.exp(logw) * x_mask * length_scale
|
||||
|
||||
w_ceil = torch.ceil(w)
|
||||
y_lengths = torch.clamp_min(torch.sum(w_ceil, [1, 2]), 1).long()
|
||||
y_mask = torch.unsqueeze(commons.create_length_mask(y_lengths, None), 1).to(
|
||||
x_mask.dtype
|
||||
)
|
||||
attn_mask = torch.unsqueeze(x_mask, 2) * torch.unsqueeze(y_mask, -1)
|
||||
attn = commons.compute_alignment_path(w_ceil, attn_mask)
|
||||
|
||||
m_p = torch.matmul(attn.squeeze(1), m_p.transpose(1, 2)).transpose(
|
||||
1, 2
|
||||
)
|
||||
logs_p = torch.matmul(attn.squeeze(1), logs_p.transpose(1, 2)).transpose(
|
||||
1, 2
|
||||
)
|
||||
|
||||
z_p = m_p + torch.randn_like(m_p) * torch.exp(logs_p) * noise_scale
|
||||
z = self.flow(z_p, y_mask, g=g, reverse=True)
|
||||
o = self.dec((z * y_mask)[:, :, :max_len], g=g)
|
||||
return o, attn, y_mask, (z, z_p, m_p, logs_p)
|
||||
@@ -0,0 +1 @@
|
||||
# Neural network building blocks
|
||||
@@ -0,0 +1,424 @@
|
||||
import math
|
||||
import torch
|
||||
from torch import nn
|
||||
from torch.nn import functional as F
|
||||
|
||||
from . import commons
|
||||
import logging
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class ChannelLayerNorm(nn.Module):
|
||||
def __init__(self, channels, eps=1e-5):
|
||||
super().__init__()
|
||||
self.channels = channels
|
||||
self.eps = eps
|
||||
|
||||
self.gamma = nn.Parameter(torch.ones(channels))
|
||||
self.beta = nn.Parameter(torch.zeros(channels))
|
||||
|
||||
def forward(self, x):
|
||||
x = x.transpose(1, -1)
|
||||
x = F.layer_norm(x, (self.channels,), self.gamma, self.beta, self.eps)
|
||||
return x.transpose(1, -1)
|
||||
|
||||
|
||||
@torch.jit.script
|
||||
def gated_activation(input_a, input_b, n_channels):
|
||||
n_channels_int = n_channels[0]
|
||||
in_act = input_a + input_b
|
||||
t_act = torch.tanh(in_act[:, :n_channels_int, :])
|
||||
s_act = torch.sigmoid(in_act[:, n_channels_int:, :])
|
||||
acts = t_act * s_act
|
||||
return acts
|
||||
|
||||
|
||||
class TransformerBlock(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
hidden_channels,
|
||||
filter_channels,
|
||||
n_heads,
|
||||
n_layers,
|
||||
kernel_size=1,
|
||||
p_dropout=0.0,
|
||||
window_size=4,
|
||||
isflow=True,
|
||||
**kwargs
|
||||
):
|
||||
super().__init__()
|
||||
self.hidden_channels = hidden_channels
|
||||
self.filter_channels = filter_channels
|
||||
self.n_heads = n_heads
|
||||
self.n_layers = n_layers
|
||||
self.kernel_size = kernel_size
|
||||
self.p_dropout = p_dropout
|
||||
self.window_size = window_size
|
||||
|
||||
self.cond_layer_idx = self.n_layers
|
||||
if "gin_channels" in kwargs:
|
||||
self.gin_channels = kwargs["gin_channels"]
|
||||
if self.gin_channels != 0:
|
||||
self.spk_emb_linear = nn.Linear(self.gin_channels, self.hidden_channels)
|
||||
self.cond_layer_idx = (
|
||||
kwargs["cond_layer_idx"] if "cond_layer_idx" in kwargs else 2
|
||||
)
|
||||
assert (
|
||||
self.cond_layer_idx < self.n_layers
|
||||
), "cond_layer_idx should be less than n_layers"
|
||||
self.drop = nn.Dropout(p_dropout)
|
||||
self.attn_layers = nn.ModuleList()
|
||||
self.norm_layers_1 = nn.ModuleList()
|
||||
self.ffn_layers = nn.ModuleList()
|
||||
self.norm_layers_2 = nn.ModuleList()
|
||||
|
||||
for i in range(self.n_layers):
|
||||
self.attn_layers.append(
|
||||
MultiHeadSelfAttention(
|
||||
hidden_channels,
|
||||
hidden_channels,
|
||||
n_heads,
|
||||
p_dropout=p_dropout,
|
||||
window_size=window_size,
|
||||
)
|
||||
)
|
||||
self.norm_layers_1.append(ChannelLayerNorm(hidden_channels))
|
||||
self.ffn_layers.append(
|
||||
FeedForward(
|
||||
hidden_channels,
|
||||
hidden_channels,
|
||||
filter_channels,
|
||||
kernel_size,
|
||||
p_dropout=p_dropout,
|
||||
)
|
||||
)
|
||||
self.norm_layers_2.append(ChannelLayerNorm(hidden_channels))
|
||||
|
||||
def forward(self, x, x_mask, g=None):
|
||||
attn_mask = x_mask.unsqueeze(2) * x_mask.unsqueeze(-1)
|
||||
x = x * x_mask
|
||||
for i in range(self.n_layers):
|
||||
if i == self.cond_layer_idx and g is not None:
|
||||
g = self.spk_emb_linear(g.transpose(1, 2))
|
||||
g = g.transpose(1, 2)
|
||||
x = x + g
|
||||
x = x * x_mask
|
||||
y = self.attn_layers[i](x, x, attn_mask)
|
||||
y = self.drop(y)
|
||||
x = self.norm_layers_1[i](x + y)
|
||||
|
||||
y = self.ffn_layers[i](x, x_mask)
|
||||
y = self.drop(y)
|
||||
x = self.norm_layers_2[i](x + y)
|
||||
x = x * x_mask
|
||||
return x
|
||||
|
||||
|
||||
class TransformerDecoder(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
hidden_channels,
|
||||
filter_channels,
|
||||
n_heads,
|
||||
n_layers,
|
||||
kernel_size=1,
|
||||
p_dropout=0.0,
|
||||
proximal_bias=False,
|
||||
proximal_init=True,
|
||||
**kwargs
|
||||
):
|
||||
super().__init__()
|
||||
self.hidden_channels = hidden_channels
|
||||
self.filter_channels = filter_channels
|
||||
self.n_heads = n_heads
|
||||
self.n_layers = n_layers
|
||||
self.kernel_size = kernel_size
|
||||
self.p_dropout = p_dropout
|
||||
self.proximal_bias = proximal_bias
|
||||
self.proximal_init = proximal_init
|
||||
|
||||
self.drop = nn.Dropout(p_dropout)
|
||||
self.self_attn_layers = nn.ModuleList()
|
||||
self.norm_layers_0 = nn.ModuleList()
|
||||
self.encdec_attn_layers = nn.ModuleList()
|
||||
self.norm_layers_1 = nn.ModuleList()
|
||||
self.ffn_layers = nn.ModuleList()
|
||||
self.norm_layers_2 = nn.ModuleList()
|
||||
for i in range(self.n_layers):
|
||||
self.self_attn_layers.append(
|
||||
MultiHeadSelfAttention(
|
||||
hidden_channels,
|
||||
hidden_channels,
|
||||
n_heads,
|
||||
p_dropout=p_dropout,
|
||||
proximal_bias=proximal_bias,
|
||||
proximal_init=proximal_init,
|
||||
)
|
||||
)
|
||||
self.norm_layers_0.append(ChannelLayerNorm(hidden_channels))
|
||||
self.encdec_attn_layers.append(
|
||||
MultiHeadSelfAttention(
|
||||
hidden_channels, hidden_channels, n_heads, p_dropout=p_dropout
|
||||
)
|
||||
)
|
||||
self.norm_layers_1.append(ChannelLayerNorm(hidden_channels))
|
||||
self.ffn_layers.append(
|
||||
FeedForward(
|
||||
hidden_channels,
|
||||
hidden_channels,
|
||||
filter_channels,
|
||||
kernel_size,
|
||||
p_dropout=p_dropout,
|
||||
causal=True,
|
||||
)
|
||||
)
|
||||
self.norm_layers_2.append(ChannelLayerNorm(hidden_channels))
|
||||
|
||||
def forward(self, x, x_mask, h, h_mask):
|
||||
self_attn_mask = commons.subsequent_mask(x_mask.size(2)).to(
|
||||
device=x.device, dtype=x.dtype
|
||||
)
|
||||
encdec_attn_mask = h_mask.unsqueeze(2) * x_mask.unsqueeze(-1)
|
||||
x = x * x_mask
|
||||
for i in range(self.n_layers):
|
||||
y = self.self_attn_layers[i](x, x, self_attn_mask)
|
||||
y = self.drop(y)
|
||||
x = self.norm_layers_0[i](x + y)
|
||||
|
||||
y = self.encdec_attn_layers[i](x, h, encdec_attn_mask)
|
||||
y = self.drop(y)
|
||||
x = self.norm_layers_1[i](x + y)
|
||||
|
||||
y = self.ffn_layers[i](x, x_mask)
|
||||
y = self.drop(y)
|
||||
x = self.norm_layers_2[i](x + y)
|
||||
x = x * x_mask
|
||||
return x
|
||||
|
||||
|
||||
class MultiHeadSelfAttention(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
channels,
|
||||
out_channels,
|
||||
n_heads,
|
||||
p_dropout=0.0,
|
||||
window_size=None,
|
||||
heads_share=True,
|
||||
block_length=None,
|
||||
proximal_bias=False,
|
||||
proximal_init=False,
|
||||
):
|
||||
super().__init__()
|
||||
assert channels % n_heads == 0
|
||||
|
||||
self.channels = channels
|
||||
self.out_channels = out_channels
|
||||
self.n_heads = n_heads
|
||||
self.p_dropout = p_dropout
|
||||
self.window_size = window_size
|
||||
self.heads_share = heads_share
|
||||
self.block_length = block_length
|
||||
self.proximal_bias = proximal_bias
|
||||
self.proximal_init = proximal_init
|
||||
self.attn = None
|
||||
|
||||
self.k_channels = channels // n_heads
|
||||
self.conv_q = nn.Conv1d(channels, channels, 1)
|
||||
self.conv_k = nn.Conv1d(channels, channels, 1)
|
||||
self.conv_v = nn.Conv1d(channels, channels, 1)
|
||||
self.conv_o = nn.Conv1d(channels, out_channels, 1)
|
||||
self.drop = nn.Dropout(p_dropout)
|
||||
|
||||
if window_size is not None:
|
||||
n_heads_rel = 1 if heads_share else n_heads
|
||||
rel_stddev = self.k_channels**-0.5
|
||||
self.emb_rel_k = nn.Parameter(
|
||||
torch.randn(n_heads_rel, window_size * 2 + 1, self.k_channels)
|
||||
* rel_stddev
|
||||
)
|
||||
self.emb_rel_v = nn.Parameter(
|
||||
torch.randn(n_heads_rel, window_size * 2 + 1, self.k_channels)
|
||||
* rel_stddev
|
||||
)
|
||||
|
||||
nn.init.xavier_uniform_(self.conv_q.weight)
|
||||
nn.init.xavier_uniform_(self.conv_k.weight)
|
||||
nn.init.xavier_uniform_(self.conv_v.weight)
|
||||
if proximal_init:
|
||||
with torch.no_grad():
|
||||
self.conv_k.weight.copy_(self.conv_q.weight)
|
||||
self.conv_k.bias.copy_(self.conv_q.bias)
|
||||
|
||||
def forward(self, x, c, attn_mask=None):
|
||||
q = self.conv_q(x)
|
||||
k = self.conv_k(c)
|
||||
v = self.conv_v(c)
|
||||
|
||||
x, self.attn = self.attention(q, k, v, mask=attn_mask)
|
||||
|
||||
x = self.conv_o(x)
|
||||
return x
|
||||
|
||||
def attention(self, query, key, value, mask=None):
|
||||
b, d, t_s, t_t = (*key.size(), query.size(2))
|
||||
query = query.view(b, self.n_heads, self.k_channels, t_t).transpose(2, 3)
|
||||
key = key.view(b, self.n_heads, self.k_channels, t_s).transpose(2, 3)
|
||||
value = value.view(b, self.n_heads, self.k_channels, t_s).transpose(2, 3)
|
||||
|
||||
scores = torch.matmul(query / math.sqrt(self.k_channels), key.transpose(-2, -1))
|
||||
if self.window_size is not None:
|
||||
assert (
|
||||
t_s == t_t
|
||||
), "Relative attention is only available for self-attention."
|
||||
key_relative_embeddings = self._get_relative_embeddings(self.emb_rel_k, t_s)
|
||||
rel_logits = self._matmul_with_relative_keys(
|
||||
query / math.sqrt(self.k_channels), key_relative_embeddings
|
||||
)
|
||||
scores_local = self._relative_position_to_absolute_position(rel_logits)
|
||||
scores = scores + scores_local
|
||||
if self.proximal_bias:
|
||||
assert t_s == t_t, "Proximal bias is only available for self-attention."
|
||||
scores = scores + self._attention_bias_proximal(t_s).to(
|
||||
device=scores.device, dtype=scores.dtype
|
||||
)
|
||||
if mask is not None:
|
||||
scores = scores.masked_fill(mask == 0, -1e4)
|
||||
if self.block_length is not None:
|
||||
assert (
|
||||
t_s == t_t
|
||||
), "Local attention is only available for self-attention."
|
||||
block_mask = (
|
||||
torch.ones_like(scores)
|
||||
.triu(-self.block_length)
|
||||
.tril(self.block_length)
|
||||
)
|
||||
scores = scores.masked_fill(block_mask == 0, -1e4)
|
||||
p_attn = F.softmax(scores, dim=-1)
|
||||
p_attn = self.drop(p_attn)
|
||||
output = torch.matmul(p_attn, value)
|
||||
if self.window_size is not None:
|
||||
relative_weights = self._absolute_position_to_relative_position(p_attn)
|
||||
value_relative_embeddings = self._get_relative_embeddings(
|
||||
self.emb_rel_v, t_s
|
||||
)
|
||||
output = output + self._matmul_with_relative_values(
|
||||
relative_weights, value_relative_embeddings
|
||||
)
|
||||
output = (
|
||||
output.transpose(2, 3).contiguous().view(b, d, t_t)
|
||||
)
|
||||
return output, p_attn
|
||||
|
||||
def _matmul_with_relative_values(self, x, y):
|
||||
ret = torch.matmul(x, y.unsqueeze(0))
|
||||
return ret
|
||||
|
||||
def _matmul_with_relative_keys(self, x, y):
|
||||
ret = torch.matmul(x, y.unsqueeze(0).transpose(-2, -1))
|
||||
return ret
|
||||
|
||||
def _get_relative_embeddings(self, relative_embeddings, length):
|
||||
2 * self.window_size + 1
|
||||
pad_length = max(length - (self.window_size + 1), 0)
|
||||
slice_start_position = max((self.window_size + 1) - length, 0)
|
||||
slice_end_position = slice_start_position + 2 * length - 1
|
||||
if pad_length > 0:
|
||||
padded_relative_embeddings = F.pad(
|
||||
relative_embeddings,
|
||||
commons.flatten_pad_shape([[0, 0], [pad_length, pad_length], [0, 0]]),
|
||||
)
|
||||
else:
|
||||
padded_relative_embeddings = relative_embeddings
|
||||
used_relative_embeddings = padded_relative_embeddings[
|
||||
:, slice_start_position:slice_end_position
|
||||
]
|
||||
return used_relative_embeddings
|
||||
|
||||
def _relative_position_to_absolute_position(self, x):
|
||||
batch, heads, length, _ = x.size()
|
||||
x = F.pad(x, commons.flatten_pad_shape([[0, 0], [0, 0], [0, 0], [0, 1]]))
|
||||
|
||||
x_flat = x.view([batch, heads, length * 2 * length])
|
||||
x_flat = F.pad(
|
||||
x_flat, commons.flatten_pad_shape([[0, 0], [0, 0], [0, length - 1]])
|
||||
)
|
||||
|
||||
x_final = x_flat.view([batch, heads, length + 1, 2 * length - 1])[
|
||||
:, :, :length, length - 1 :
|
||||
]
|
||||
return x_final
|
||||
|
||||
def _absolute_position_to_relative_position(self, x):
|
||||
batch, heads, length, _ = x.size()
|
||||
x = F.pad(
|
||||
x, commons.flatten_pad_shape([[0, 0], [0, 0], [0, 0], [0, length - 1]])
|
||||
)
|
||||
x_flat = x.view([batch, heads, length**2 + length * (length - 1)])
|
||||
x_flat = F.pad(x_flat, commons.flatten_pad_shape([[0, 0], [0, 0], [length, 0]]))
|
||||
x_final = x_flat.view([batch, heads, length, 2 * length])[:, :, :, 1:]
|
||||
return x_final
|
||||
|
||||
def _attention_bias_proximal(self, length):
|
||||
r = torch.arange(length, dtype=torch.float32)
|
||||
diff = torch.unsqueeze(r, 0) - torch.unsqueeze(r, 1)
|
||||
return torch.unsqueeze(torch.unsqueeze(-torch.log1p(torch.abs(diff)), 0), 0)
|
||||
|
||||
|
||||
class FeedForward(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
in_channels,
|
||||
out_channels,
|
||||
filter_channels,
|
||||
kernel_size,
|
||||
p_dropout=0.0,
|
||||
activation=None,
|
||||
causal=False,
|
||||
):
|
||||
super().__init__()
|
||||
self.in_channels = in_channels
|
||||
self.out_channels = out_channels
|
||||
self.filter_channels = filter_channels
|
||||
self.kernel_size = kernel_size
|
||||
self.p_dropout = p_dropout
|
||||
self.activation = activation
|
||||
self.causal = causal
|
||||
|
||||
if causal:
|
||||
self.padding = self._causal_padding
|
||||
else:
|
||||
self.padding = self._same_padding
|
||||
|
||||
self.conv_1 = nn.Conv1d(in_channels, filter_channels, kernel_size)
|
||||
self.conv_2 = nn.Conv1d(filter_channels, out_channels, kernel_size)
|
||||
self.drop = nn.Dropout(p_dropout)
|
||||
|
||||
def forward(self, x, x_mask):
|
||||
x = self.conv_1(self.padding(x * x_mask))
|
||||
if self.activation == "gelu":
|
||||
x = x * torch.sigmoid(1.702 * x)
|
||||
else:
|
||||
x = torch.relu(x)
|
||||
x = self.drop(x)
|
||||
x = self.conv_2(self.padding(x * x_mask))
|
||||
return x * x_mask
|
||||
|
||||
def _causal_padding(self, x):
|
||||
if self.kernel_size == 1:
|
||||
return x
|
||||
pad_l = self.kernel_size - 1
|
||||
pad_r = 0
|
||||
padding = [[0, 0], [0, 0], [pad_l, pad_r]]
|
||||
x = F.pad(x, commons.flatten_pad_shape(padding))
|
||||
return x
|
||||
|
||||
def _same_padding(self, x):
|
||||
if self.kernel_size == 1:
|
||||
return x
|
||||
pad_l = (self.kernel_size - 1) // 2
|
||||
pad_r = self.kernel_size // 2
|
||||
padding = [[0, 0], [0, 0], [pad_l, pad_r]]
|
||||
x = F.pad(x, commons.flatten_pad_shape(padding))
|
||||
return x
|
||||
@@ -0,0 +1,151 @@
|
||||
import math
|
||||
import torch
|
||||
from torch.nn import functional as F
|
||||
|
||||
|
||||
def initialize_weights(m, mean=0.0, std=0.01):
|
||||
classname = m.__class__.__name__
|
||||
if classname.find("Conv") != -1:
|
||||
m.weight.data.normal_(mean, std)
|
||||
|
||||
|
||||
def compute_padding(kernel_size, dilation=1):
|
||||
return int((kernel_size * dilation - dilation) / 2)
|
||||
|
||||
|
||||
def flatten_pad_shape(pad_shape):
|
||||
layer = pad_shape[::-1]
|
||||
pad_shape = [item for sublist in layer for item in sublist]
|
||||
return pad_shape
|
||||
|
||||
|
||||
def insert_blanks(lst, item):
|
||||
result = [item] * (len(lst) * 2 + 1)
|
||||
result[1::2] = lst
|
||||
return result
|
||||
|
||||
|
||||
def kl_divergence(m_p, logs_p, m_q, logs_q):
|
||||
kl = (logs_q - logs_p) - 0.5
|
||||
kl += (
|
||||
0.5 * (torch.exp(2.0 * logs_p) + ((m_p - m_q) ** 2)) * torch.exp(-2.0 * logs_q)
|
||||
)
|
||||
return kl
|
||||
|
||||
|
||||
def rand_gumbel(shape):
|
||||
uniform_samples = torch.rand(shape) * 0.99998 + 0.00001
|
||||
return -torch.log(-torch.log(uniform_samples))
|
||||
|
||||
|
||||
def rand_gumbel_like(x):
|
||||
g = rand_gumbel(x.size()).to(dtype=x.dtype, device=x.device)
|
||||
return g
|
||||
|
||||
|
||||
def extract_segments(x, ids_str, segment_size=4):
|
||||
ret = torch.zeros_like(x[:, :, :segment_size])
|
||||
for i in range(x.size(0)):
|
||||
idx_str = max(0, ids_str[i].item())
|
||||
idx_end = idx_str + segment_size
|
||||
available = x.size(2) - idx_str
|
||||
if available >= segment_size:
|
||||
ret[i] = x[i, :, idx_str:idx_end]
|
||||
elif available > 0:
|
||||
ret[i, :, :available] = x[i, :, idx_str:idx_str + available]
|
||||
return ret
|
||||
|
||||
|
||||
def random_segments(x, x_lengths=None, segment_size=4):
|
||||
b, d, t = x.size()
|
||||
if x_lengths is None:
|
||||
x_lengths = t
|
||||
ids_str_max = torch.clamp(x_lengths - segment_size + 1, min=0)
|
||||
ids_str = (torch.rand([b]).to(device=x.device) * ids_str_max).to(dtype=torch.long)
|
||||
ret = extract_segments(x, ids_str, segment_size)
|
||||
return ret, ids_str
|
||||
|
||||
|
||||
def get_timing_signal_1d(length, channels, min_timescale=1.0, max_timescale=1.0e4):
|
||||
position = torch.arange(length, dtype=torch.float)
|
||||
num_timescales = channels // 2
|
||||
log_timescale_increment = math.log(float(max_timescale) / float(min_timescale)) / (
|
||||
num_timescales - 1
|
||||
)
|
||||
inv_timescales = min_timescale * torch.exp(
|
||||
torch.arange(num_timescales, dtype=torch.float) * -log_timescale_increment
|
||||
)
|
||||
scaled_time = position.unsqueeze(0) * inv_timescales.unsqueeze(1)
|
||||
signal = torch.cat([torch.sin(scaled_time), torch.cos(scaled_time)], 0)
|
||||
signal = F.pad(signal, [0, 0, 0, channels % 2])
|
||||
signal = signal.view(1, channels, length)
|
||||
return signal
|
||||
|
||||
|
||||
def add_timing_signal_1d(x, min_timescale=1.0, max_timescale=1.0e4):
|
||||
b, channels, length = x.size()
|
||||
signal = get_timing_signal_1d(length, channels, min_timescale, max_timescale)
|
||||
return x + signal.to(dtype=x.dtype, device=x.device)
|
||||
|
||||
|
||||
def cat_timing_signal_1d(x, min_timescale=1.0, max_timescale=1.0e4, axis=1):
|
||||
b, channels, length = x.size()
|
||||
signal = get_timing_signal_1d(length, channels, min_timescale, max_timescale)
|
||||
return torch.cat([x, signal.to(dtype=x.dtype, device=x.device)], axis)
|
||||
|
||||
|
||||
def subsequent_mask(length):
|
||||
mask = torch.tril(torch.ones(length, length)).unsqueeze(0).unsqueeze(0)
|
||||
return mask
|
||||
|
||||
|
||||
@torch.jit.script
|
||||
def gated_activation(input_a, input_b, n_channels):
|
||||
n_channels_int = n_channels[0]
|
||||
in_act = input_a + input_b
|
||||
t_act = torch.tanh(in_act[:, :n_channels_int, :])
|
||||
s_act = torch.sigmoid(in_act[:, n_channels_int:, :])
|
||||
acts = t_act * s_act
|
||||
return acts
|
||||
|
||||
|
||||
def shift_1d(x):
|
||||
x = F.pad(x, flatten_pad_shape([[0, 0], [0, 0], [1, 0]]))[:, :, :-1]
|
||||
return x
|
||||
|
||||
|
||||
def create_length_mask(length, max_length=None):
|
||||
if max_length is None:
|
||||
max_length = length.max()
|
||||
x = torch.arange(max_length, dtype=length.dtype, device=length.device)
|
||||
return x.unsqueeze(0) < length.unsqueeze(1)
|
||||
|
||||
|
||||
def compute_alignment_path(duration, mask):
|
||||
b, _, t_y, t_x = mask.shape
|
||||
cum_duration = torch.cumsum(duration, -1)
|
||||
|
||||
cum_duration_flat = cum_duration.view(b * t_x)
|
||||
path = create_length_mask(cum_duration_flat, t_y).to(mask.dtype)
|
||||
path = path.view(b, t_x, t_y)
|
||||
path = path - F.pad(path, flatten_pad_shape([[0, 0], [1, 0], [0, 0]]))[:, :-1]
|
||||
path = path.unsqueeze(1).transpose(2, 3) * mask
|
||||
return path
|
||||
|
||||
|
||||
def clip_grad_value_(parameters, clip_value, norm_type=2):
|
||||
if isinstance(parameters, torch.Tensor):
|
||||
parameters = [parameters]
|
||||
parameters = list(filter(lambda p: p.grad is not None, parameters))
|
||||
norm_type = float(norm_type)
|
||||
if clip_value is not None:
|
||||
clip_value = float(clip_value)
|
||||
|
||||
total_norm = 0
|
||||
for p in parameters:
|
||||
param_norm = p.grad.data.norm(norm_type)
|
||||
total_norm += param_norm.item() ** norm_type
|
||||
if clip_value is not None:
|
||||
p.grad.data.clamp_(min=-clip_value, max=clip_value)
|
||||
total_norm = total_norm ** (1.0 / norm_type)
|
||||
return total_norm
|
||||
@@ -0,0 +1,578 @@
|
||||
import math
|
||||
import torch
|
||||
from torch import nn
|
||||
from torch.nn import functional as F
|
||||
|
||||
from torch.nn import Conv1d
|
||||
from torch.nn.utils import weight_norm, remove_weight_norm
|
||||
|
||||
from . import commons
|
||||
from .commons import initialize_weights, compute_padding
|
||||
from .transforms import spline_transform
|
||||
from .attentions import TransformerBlock
|
||||
|
||||
LRELU_SLOPE = 0.1
|
||||
|
||||
|
||||
class ChannelNorm(nn.Module):
|
||||
def __init__(self, channels, eps=1e-5):
|
||||
super().__init__()
|
||||
self.channels = channels
|
||||
self.eps = eps
|
||||
|
||||
self.gamma = nn.Parameter(torch.ones(channels))
|
||||
self.beta = nn.Parameter(torch.zeros(channels))
|
||||
|
||||
def forward(self, x):
|
||||
x = x.transpose(1, -1)
|
||||
x = F.layer_norm(x, (self.channels,), self.gamma, self.beta, self.eps)
|
||||
return x.transpose(1, -1)
|
||||
|
||||
|
||||
class ConvReluNorm(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
in_channels,
|
||||
hidden_channels,
|
||||
out_channels,
|
||||
kernel_size,
|
||||
n_layers,
|
||||
p_dropout,
|
||||
):
|
||||
super().__init__()
|
||||
self.in_channels = in_channels
|
||||
self.hidden_channels = hidden_channels
|
||||
self.out_channels = out_channels
|
||||
self.kernel_size = kernel_size
|
||||
self.n_layers = n_layers
|
||||
self.p_dropout = p_dropout
|
||||
assert n_layers > 1, "Number of layers should be larger than 0."
|
||||
|
||||
self.conv_layers = nn.ModuleList()
|
||||
self.norm_layers = nn.ModuleList()
|
||||
self.conv_layers.append(
|
||||
nn.Conv1d(
|
||||
in_channels, hidden_channels, kernel_size, padding=kernel_size // 2
|
||||
)
|
||||
)
|
||||
self.norm_layers.append(ChannelNorm(hidden_channels))
|
||||
self.relu_drop = nn.Sequential(nn.ReLU(), nn.Dropout(p_dropout))
|
||||
for _ in range(n_layers - 1):
|
||||
self.conv_layers.append(
|
||||
nn.Conv1d(
|
||||
hidden_channels,
|
||||
hidden_channels,
|
||||
kernel_size,
|
||||
padding=kernel_size // 2,
|
||||
)
|
||||
)
|
||||
self.norm_layers.append(ChannelNorm(hidden_channels))
|
||||
self.proj = nn.Conv1d(hidden_channels, out_channels, 1)
|
||||
self.proj.weight.data.zero_()
|
||||
self.proj.bias.data.zero_()
|
||||
|
||||
def forward(self, x, x_mask):
|
||||
x_org = x
|
||||
for i in range(self.n_layers):
|
||||
x = self.conv_layers[i](x * x_mask)
|
||||
x = self.norm_layers[i](x)
|
||||
x = self.relu_drop(x)
|
||||
x = x_org + self.proj(x)
|
||||
return x * x_mask
|
||||
|
||||
|
||||
class DepthwiseSepConv(nn.Module):
|
||||
"""Dilated and Depth-Separable Convolution"""
|
||||
|
||||
def __init__(self, channels, kernel_size, n_layers, p_dropout=0.0):
|
||||
super().__init__()
|
||||
self.channels = channels
|
||||
self.kernel_size = kernel_size
|
||||
self.n_layers = n_layers
|
||||
self.p_dropout = p_dropout
|
||||
|
||||
self.drop = nn.Dropout(p_dropout)
|
||||
self.convs_sep = nn.ModuleList()
|
||||
self.convs_1x1 = nn.ModuleList()
|
||||
self.norms_1 = nn.ModuleList()
|
||||
self.norms_2 = nn.ModuleList()
|
||||
for i in range(n_layers):
|
||||
dilation = kernel_size**i
|
||||
padding = (kernel_size * dilation - dilation) // 2
|
||||
self.convs_sep.append(
|
||||
nn.Conv1d(
|
||||
channels,
|
||||
channels,
|
||||
kernel_size,
|
||||
groups=channels,
|
||||
dilation=dilation,
|
||||
padding=padding,
|
||||
)
|
||||
)
|
||||
self.convs_1x1.append(nn.Conv1d(channels, channels, 1))
|
||||
self.norms_1.append(ChannelNorm(channels))
|
||||
self.norms_2.append(ChannelNorm(channels))
|
||||
|
||||
def forward(self, x, x_mask, g=None):
|
||||
if g is not None:
|
||||
x = x + g
|
||||
for i in range(self.n_layers):
|
||||
y = self.convs_sep[i](x * x_mask)
|
||||
y = self.norms_1[i](y)
|
||||
y = F.gelu(y)
|
||||
y = self.convs_1x1[i](y)
|
||||
y = self.norms_2[i](y)
|
||||
y = F.gelu(y)
|
||||
y = self.drop(y)
|
||||
x = x + y
|
||||
return x * x_mask
|
||||
|
||||
|
||||
class WaveNet(torch.nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
hidden_channels,
|
||||
kernel_size,
|
||||
dilation_rate,
|
||||
n_layers,
|
||||
gin_channels=0,
|
||||
p_dropout=0,
|
||||
):
|
||||
super(WaveNet, self).__init__()
|
||||
assert kernel_size % 2 == 1
|
||||
self.hidden_channels = hidden_channels
|
||||
self.kernel_size = (kernel_size,)
|
||||
self.dilation_rate = dilation_rate
|
||||
self.n_layers = n_layers
|
||||
self.gin_channels = gin_channels
|
||||
self.p_dropout = p_dropout
|
||||
|
||||
self.in_layers = torch.nn.ModuleList()
|
||||
self.res_skip_layers = torch.nn.ModuleList()
|
||||
self.drop = nn.Dropout(p_dropout)
|
||||
|
||||
if gin_channels != 0:
|
||||
cond_layer = torch.nn.Conv1d(
|
||||
gin_channels, 2 * hidden_channels * n_layers, 1
|
||||
)
|
||||
self.cond_layer = torch.nn.utils.weight_norm(cond_layer, name="weight")
|
||||
|
||||
for i in range(n_layers):
|
||||
dilation = dilation_rate**i
|
||||
padding = int((kernel_size * dilation - dilation) / 2)
|
||||
in_layer = torch.nn.Conv1d(
|
||||
hidden_channels,
|
||||
2 * hidden_channels,
|
||||
kernel_size,
|
||||
dilation=dilation,
|
||||
padding=padding,
|
||||
)
|
||||
in_layer = torch.nn.utils.weight_norm(in_layer, name="weight")
|
||||
self.in_layers.append(in_layer)
|
||||
|
||||
if i < n_layers - 1:
|
||||
res_skip_channels = 2 * hidden_channels
|
||||
else:
|
||||
res_skip_channels = hidden_channels
|
||||
|
||||
res_skip_layer = torch.nn.Conv1d(hidden_channels, res_skip_channels, 1)
|
||||
res_skip_layer = torch.nn.utils.weight_norm(res_skip_layer, name="weight")
|
||||
self.res_skip_layers.append(res_skip_layer)
|
||||
|
||||
def forward(self, x, x_mask, g=None, **kwargs):
|
||||
output = torch.zeros_like(x)
|
||||
n_channels_tensor = torch.IntTensor([self.hidden_channels])
|
||||
|
||||
if g is not None:
|
||||
g = self.cond_layer(g)
|
||||
|
||||
for i in range(self.n_layers):
|
||||
x_in = self.in_layers[i](x)
|
||||
if g is not None:
|
||||
cond_offset = i * 2 * self.hidden_channels
|
||||
g_l = g[:, cond_offset : cond_offset + 2 * self.hidden_channels, :]
|
||||
else:
|
||||
g_l = torch.zeros_like(x_in)
|
||||
|
||||
acts = commons.gated_activation(x_in, g_l, n_channels_tensor)
|
||||
acts = self.drop(acts)
|
||||
|
||||
res_skip_acts = self.res_skip_layers[i](acts)
|
||||
if i < self.n_layers - 1:
|
||||
res_acts = res_skip_acts[:, : self.hidden_channels, :]
|
||||
x = (x + res_acts) * x_mask
|
||||
output = output + res_skip_acts[:, self.hidden_channels :, :]
|
||||
else:
|
||||
output = output + res_skip_acts
|
||||
return output * x_mask
|
||||
|
||||
def remove_weight_norm(self):
|
||||
if self.gin_channels != 0:
|
||||
torch.nn.utils.remove_weight_norm(self.cond_layer)
|
||||
for l in self.in_layers:
|
||||
torch.nn.utils.remove_weight_norm(l)
|
||||
for l in self.res_skip_layers:
|
||||
torch.nn.utils.remove_weight_norm(l)
|
||||
|
||||
|
||||
class ConvResBlock(torch.nn.Module):
|
||||
def __init__(self, channels, kernel_size=3, dilation=(1, 3, 5)):
|
||||
super(ConvResBlock, self).__init__()
|
||||
self.convs1 = nn.ModuleList(
|
||||
[
|
||||
weight_norm(
|
||||
Conv1d(
|
||||
channels,
|
||||
channels,
|
||||
kernel_size,
|
||||
1,
|
||||
dilation=dilation[0],
|
||||
padding=compute_padding(kernel_size, dilation[0]),
|
||||
)
|
||||
),
|
||||
weight_norm(
|
||||
Conv1d(
|
||||
channels,
|
||||
channels,
|
||||
kernel_size,
|
||||
1,
|
||||
dilation=dilation[1],
|
||||
padding=compute_padding(kernel_size, dilation[1]),
|
||||
)
|
||||
),
|
||||
weight_norm(
|
||||
Conv1d(
|
||||
channels,
|
||||
channels,
|
||||
kernel_size,
|
||||
1,
|
||||
dilation=dilation[2],
|
||||
padding=compute_padding(kernel_size, dilation[2]),
|
||||
)
|
||||
),
|
||||
]
|
||||
)
|
||||
self.convs1.apply(initialize_weights)
|
||||
|
||||
self.convs2 = nn.ModuleList(
|
||||
[
|
||||
weight_norm(
|
||||
Conv1d(
|
||||
channels,
|
||||
channels,
|
||||
kernel_size,
|
||||
1,
|
||||
dilation=1,
|
||||
padding=compute_padding(kernel_size, 1),
|
||||
)
|
||||
),
|
||||
weight_norm(
|
||||
Conv1d(
|
||||
channels,
|
||||
channels,
|
||||
kernel_size,
|
||||
1,
|
||||
dilation=1,
|
||||
padding=compute_padding(kernel_size, 1),
|
||||
)
|
||||
),
|
||||
weight_norm(
|
||||
Conv1d(
|
||||
channels,
|
||||
channels,
|
||||
kernel_size,
|
||||
1,
|
||||
dilation=1,
|
||||
padding=compute_padding(kernel_size, 1),
|
||||
)
|
||||
),
|
||||
]
|
||||
)
|
||||
self.convs2.apply(initialize_weights)
|
||||
|
||||
def forward(self, x, x_mask=None):
|
||||
for c1, c2 in zip(self.convs1, self.convs2):
|
||||
xt = F.leaky_relu(x, LRELU_SLOPE)
|
||||
if x_mask is not None:
|
||||
xt = xt * x_mask
|
||||
xt = c1(xt)
|
||||
xt = F.leaky_relu(xt, LRELU_SLOPE)
|
||||
if x_mask is not None:
|
||||
xt = xt * x_mask
|
||||
xt = c2(xt)
|
||||
x = xt + x
|
||||
if x_mask is not None:
|
||||
x = x * x_mask
|
||||
return x
|
||||
|
||||
def remove_weight_norm(self):
|
||||
for l in self.convs1:
|
||||
remove_weight_norm(l)
|
||||
for l in self.convs2:
|
||||
remove_weight_norm(l)
|
||||
|
||||
|
||||
class ConvResBlockLight(torch.nn.Module):
|
||||
def __init__(self, channels, kernel_size=3, dilation=(1, 3)):
|
||||
super(ConvResBlockLight, self).__init__()
|
||||
self.convs = nn.ModuleList(
|
||||
[
|
||||
weight_norm(
|
||||
Conv1d(
|
||||
channels,
|
||||
channels,
|
||||
kernel_size,
|
||||
1,
|
||||
dilation=dilation[0],
|
||||
padding=compute_padding(kernel_size, dilation[0]),
|
||||
)
|
||||
),
|
||||
weight_norm(
|
||||
Conv1d(
|
||||
channels,
|
||||
channels,
|
||||
kernel_size,
|
||||
1,
|
||||
dilation=dilation[1],
|
||||
padding=compute_padding(kernel_size, dilation[1]),
|
||||
)
|
||||
),
|
||||
]
|
||||
)
|
||||
self.convs.apply(initialize_weights)
|
||||
|
||||
def forward(self, x, x_mask=None):
|
||||
for c in self.convs:
|
||||
xt = F.leaky_relu(x, LRELU_SLOPE)
|
||||
if x_mask is not None:
|
||||
xt = xt * x_mask
|
||||
xt = c(xt)
|
||||
x = xt + x
|
||||
if x_mask is not None:
|
||||
x = x * x_mask
|
||||
return x
|
||||
|
||||
def remove_weight_norm(self):
|
||||
for l in self.convs:
|
||||
remove_weight_norm(l)
|
||||
|
||||
|
||||
class LogTransform(nn.Module):
|
||||
def forward(self, x, x_mask, reverse=False, **kwargs):
|
||||
if not reverse:
|
||||
y = torch.log(torch.clamp_min(x, 1e-5)) * x_mask
|
||||
logdet = torch.sum(-y, [1, 2])
|
||||
return y, logdet
|
||||
else:
|
||||
x = torch.exp(x) * x_mask
|
||||
return x
|
||||
|
||||
|
||||
class FlipTransform(nn.Module):
|
||||
def forward(self, x, *args, reverse=False, **kwargs):
|
||||
x = torch.flip(x, [1])
|
||||
if not reverse:
|
||||
logdet = torch.zeros(x.size(0)).to(dtype=x.dtype, device=x.device)
|
||||
return x, logdet
|
||||
else:
|
||||
return x
|
||||
|
||||
|
||||
class AffineCoupling(nn.Module):
|
||||
def __init__(self, channels):
|
||||
super().__init__()
|
||||
self.channels = channels
|
||||
self.m = nn.Parameter(torch.zeros(channels, 1))
|
||||
self.logs = nn.Parameter(torch.zeros(channels, 1))
|
||||
|
||||
def forward(self, x, x_mask, reverse=False, **kwargs):
|
||||
if not reverse:
|
||||
y = self.m + torch.exp(self.logs) * x
|
||||
y = y * x_mask
|
||||
logdet = torch.sum(self.logs * x_mask, [1, 2])
|
||||
return y, logdet
|
||||
else:
|
||||
x = (x - self.m) * torch.exp(-self.logs) * x_mask
|
||||
return x
|
||||
|
||||
|
||||
class FlowCouplingLayer(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
channels,
|
||||
hidden_channels,
|
||||
kernel_size,
|
||||
dilation_rate,
|
||||
n_layers,
|
||||
p_dropout=0,
|
||||
gin_channels=0,
|
||||
mean_only=False,
|
||||
):
|
||||
assert channels % 2 == 0, "channels should be divisible by 2"
|
||||
super().__init__()
|
||||
self.channels = channels
|
||||
self.hidden_channels = hidden_channels
|
||||
self.kernel_size = kernel_size
|
||||
self.dilation_rate = dilation_rate
|
||||
self.n_layers = n_layers
|
||||
self.half_channels = channels // 2
|
||||
self.mean_only = mean_only
|
||||
|
||||
self.pre = nn.Conv1d(self.half_channels, hidden_channels, 1)
|
||||
self.enc = WaveNet(
|
||||
hidden_channels,
|
||||
kernel_size,
|
||||
dilation_rate,
|
||||
n_layers,
|
||||
p_dropout=p_dropout,
|
||||
gin_channels=gin_channels,
|
||||
)
|
||||
self.post = nn.Conv1d(hidden_channels, self.half_channels * (2 - mean_only), 1)
|
||||
self.post.weight.data.zero_()
|
||||
self.post.bias.data.zero_()
|
||||
|
||||
def forward(self, x, x_mask, g=None, reverse=False):
|
||||
x0, x1 = torch.split(x, [self.half_channels] * 2, 1)
|
||||
h = self.pre(x0) * x_mask
|
||||
h = self.enc(h, x_mask, g=g)
|
||||
stats = self.post(h) * x_mask
|
||||
if not self.mean_only:
|
||||
m, logs = torch.split(stats, [self.half_channels] * 2, 1)
|
||||
else:
|
||||
m = stats
|
||||
logs = torch.zeros_like(m)
|
||||
|
||||
if not reverse:
|
||||
x1 = m + x1 * torch.exp(logs) * x_mask
|
||||
x = torch.cat([x0, x1], 1)
|
||||
logdet = torch.sum(logs, [1, 2])
|
||||
return x, logdet
|
||||
else:
|
||||
x1 = (x1 - m) * torch.exp(-logs) * x_mask
|
||||
x = torch.cat([x0, x1], 1)
|
||||
return x
|
||||
|
||||
|
||||
class ConvolutionalFlow(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
in_channels,
|
||||
filter_channels,
|
||||
kernel_size,
|
||||
n_layers,
|
||||
num_bins=10,
|
||||
tail_bound=5.0,
|
||||
):
|
||||
super().__init__()
|
||||
self.in_channels = in_channels
|
||||
self.filter_channels = filter_channels
|
||||
self.kernel_size = kernel_size
|
||||
self.n_layers = n_layers
|
||||
self.num_bins = num_bins
|
||||
self.tail_bound = tail_bound
|
||||
self.half_channels = in_channels // 2
|
||||
|
||||
self.pre = nn.Conv1d(self.half_channels, filter_channels, 1)
|
||||
self.convs = DepthwiseSepConv(filter_channels, kernel_size, n_layers, p_dropout=0.0)
|
||||
self.proj = nn.Conv1d(
|
||||
filter_channels, self.half_channels * (num_bins * 3 - 1), 1
|
||||
)
|
||||
self.proj.weight.data.zero_()
|
||||
self.proj.bias.data.zero_()
|
||||
|
||||
def forward(self, x, x_mask, g=None, reverse=False):
|
||||
x0, x1 = torch.split(x, [self.half_channels] * 2, 1)
|
||||
h = self.pre(x0)
|
||||
h = self.convs(h, x_mask, g=g)
|
||||
h = self.proj(h) * x_mask
|
||||
|
||||
b, c, t = x0.shape
|
||||
h = h.reshape(b, c, -1, t).permute(0, 1, 3, 2)
|
||||
|
||||
unnormalized_widths = h[..., : self.num_bins] / math.sqrt(self.filter_channels)
|
||||
unnormalized_heights = h[..., self.num_bins : 2 * self.num_bins] / math.sqrt(
|
||||
self.filter_channels
|
||||
)
|
||||
unnormalized_derivatives = h[..., 2 * self.num_bins :]
|
||||
|
||||
x1, logabsdet = spline_transform(
|
||||
x1,
|
||||
unnormalized_widths,
|
||||
unnormalized_heights,
|
||||
unnormalized_derivatives,
|
||||
inverse=reverse,
|
||||
tails="linear",
|
||||
tail_bound=self.tail_bound,
|
||||
)
|
||||
|
||||
x = torch.cat([x0, x1], 1) * x_mask
|
||||
logdet = torch.sum(logabsdet * x_mask, [1, 2])
|
||||
if not reverse:
|
||||
return x, logdet
|
||||
else:
|
||||
return x
|
||||
|
||||
|
||||
class TransformerCouplingLayer(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
channels,
|
||||
hidden_channels,
|
||||
kernel_size,
|
||||
n_layers,
|
||||
n_heads,
|
||||
p_dropout=0,
|
||||
filter_channels=0,
|
||||
mean_only=False,
|
||||
wn_sharing_parameter=None,
|
||||
gin_channels=0,
|
||||
):
|
||||
assert n_layers == 3, n_layers
|
||||
assert channels % 2 == 0, "channels should be divisible by 2"
|
||||
super().__init__()
|
||||
self.channels = channels
|
||||
self.hidden_channels = hidden_channels
|
||||
self.kernel_size = kernel_size
|
||||
self.n_layers = n_layers
|
||||
self.half_channels = channels // 2
|
||||
self.mean_only = mean_only
|
||||
|
||||
self.pre = nn.Conv1d(self.half_channels, hidden_channels, 1)
|
||||
self.enc = (
|
||||
TransformerBlock(
|
||||
hidden_channels,
|
||||
filter_channels,
|
||||
n_heads,
|
||||
n_layers,
|
||||
kernel_size,
|
||||
p_dropout,
|
||||
isflow=True,
|
||||
gin_channels=gin_channels,
|
||||
)
|
||||
if wn_sharing_parameter is None
|
||||
else wn_sharing_parameter
|
||||
)
|
||||
self.post = nn.Conv1d(hidden_channels, self.half_channels * (2 - mean_only), 1)
|
||||
self.post.weight.data.zero_()
|
||||
self.post.bias.data.zero_()
|
||||
|
||||
def forward(self, x, x_mask, g=None, reverse=False):
|
||||
x0, x1 = torch.split(x, [self.half_channels] * 2, 1)
|
||||
h = self.pre(x0) * x_mask
|
||||
h = self.enc(h, x_mask, g=g)
|
||||
stats = self.post(h) * x_mask
|
||||
if not self.mean_only:
|
||||
m, logs = torch.split(stats, [self.half_channels] * 2, 1)
|
||||
else:
|
||||
m = stats
|
||||
logs = torch.zeros_like(m)
|
||||
|
||||
if not reverse:
|
||||
x1 = m + x1 * torch.exp(logs) * x_mask
|
||||
x = torch.cat([x0, x1], 1)
|
||||
logdet = torch.sum(logs, [1, 2])
|
||||
return x, logdet
|
||||
else:
|
||||
x1 = (x1 - m) * torch.exp(-logs) * x_mask
|
||||
x = torch.cat([x0, x1], 1)
|
||||
return x
|
||||
@@ -0,0 +1,209 @@
|
||||
import torch
|
||||
from torch.nn import functional as F
|
||||
|
||||
import numpy as np
|
||||
|
||||
|
||||
DEFAULT_MIN_BIN_WIDTH = 1e-3
|
||||
DEFAULT_MIN_BIN_HEIGHT = 1e-3
|
||||
DEFAULT_MIN_DERIVATIVE = 1e-3
|
||||
|
||||
|
||||
def spline_transform(
|
||||
inputs,
|
||||
unnormalized_widths,
|
||||
unnormalized_heights,
|
||||
unnormalized_derivatives,
|
||||
inverse=False,
|
||||
tails=None,
|
||||
tail_bound=1.0,
|
||||
min_bin_width=DEFAULT_MIN_BIN_WIDTH,
|
||||
min_bin_height=DEFAULT_MIN_BIN_HEIGHT,
|
||||
min_derivative=DEFAULT_MIN_DERIVATIVE,
|
||||
):
|
||||
if tails is None:
|
||||
spline_fn = quadratic_spline
|
||||
spline_kwargs = {}
|
||||
else:
|
||||
spline_fn = unbounded_spline
|
||||
spline_kwargs = {"tails": tails, "tail_bound": tail_bound}
|
||||
|
||||
outputs, logabsdet = spline_fn(
|
||||
inputs=inputs,
|
||||
unnormalized_widths=unnormalized_widths,
|
||||
unnormalized_heights=unnormalized_heights,
|
||||
unnormalized_derivatives=unnormalized_derivatives,
|
||||
inverse=inverse,
|
||||
min_bin_width=min_bin_width,
|
||||
min_bin_height=min_bin_height,
|
||||
min_derivative=min_derivative,
|
||||
**spline_kwargs
|
||||
)
|
||||
return outputs, logabsdet
|
||||
|
||||
|
||||
def searchsorted(bin_locations, inputs, eps=1e-6):
|
||||
bin_locations[..., -1] += eps
|
||||
return torch.sum(inputs[..., None] >= bin_locations, dim=-1) - 1
|
||||
|
||||
|
||||
def unbounded_spline(
|
||||
inputs,
|
||||
unnormalized_widths,
|
||||
unnormalized_heights,
|
||||
unnormalized_derivatives,
|
||||
inverse=False,
|
||||
tails="linear",
|
||||
tail_bound=1.0,
|
||||
min_bin_width=DEFAULT_MIN_BIN_WIDTH,
|
||||
min_bin_height=DEFAULT_MIN_BIN_HEIGHT,
|
||||
min_derivative=DEFAULT_MIN_DERIVATIVE,
|
||||
):
|
||||
inside_interval_mask = (inputs >= -tail_bound) & (inputs <= tail_bound)
|
||||
outside_interval_mask = ~inside_interval_mask
|
||||
|
||||
outputs = torch.zeros_like(inputs)
|
||||
logabsdet = torch.zeros_like(inputs)
|
||||
|
||||
if tails == "linear":
|
||||
unnormalized_derivatives = F.pad(unnormalized_derivatives, pad=(1, 1))
|
||||
constant = np.log(np.exp(1 - min_derivative) - 1)
|
||||
unnormalized_derivatives[..., 0] = constant
|
||||
unnormalized_derivatives[..., -1] = constant
|
||||
|
||||
outputs[outside_interval_mask] = inputs[outside_interval_mask]
|
||||
logabsdet[outside_interval_mask] = 0
|
||||
else:
|
||||
raise RuntimeError("{} tails are not implemented.".format(tails))
|
||||
|
||||
(
|
||||
outputs[inside_interval_mask],
|
||||
logabsdet[inside_interval_mask],
|
||||
) = quadratic_spline(
|
||||
inputs=inputs[inside_interval_mask],
|
||||
unnormalized_widths=unnormalized_widths[inside_interval_mask, :],
|
||||
unnormalized_heights=unnormalized_heights[inside_interval_mask, :],
|
||||
unnormalized_derivatives=unnormalized_derivatives[inside_interval_mask, :],
|
||||
inverse=inverse,
|
||||
left=-tail_bound,
|
||||
right=tail_bound,
|
||||
bottom=-tail_bound,
|
||||
top=tail_bound,
|
||||
min_bin_width=min_bin_width,
|
||||
min_bin_height=min_bin_height,
|
||||
min_derivative=min_derivative,
|
||||
)
|
||||
|
||||
return outputs, logabsdet
|
||||
|
||||
|
||||
def quadratic_spline(
|
||||
inputs,
|
||||
unnormalized_widths,
|
||||
unnormalized_heights,
|
||||
unnormalized_derivatives,
|
||||
inverse=False,
|
||||
left=0.0,
|
||||
right=1.0,
|
||||
bottom=0.0,
|
||||
top=1.0,
|
||||
min_bin_width=DEFAULT_MIN_BIN_WIDTH,
|
||||
min_bin_height=DEFAULT_MIN_BIN_HEIGHT,
|
||||
min_derivative=DEFAULT_MIN_DERIVATIVE,
|
||||
):
|
||||
if torch.min(inputs) < left or torch.max(inputs) > right:
|
||||
raise ValueError("Input to a transform is not within its domain")
|
||||
|
||||
num_bins = unnormalized_widths.shape[-1]
|
||||
|
||||
if min_bin_width * num_bins > 1.0:
|
||||
raise ValueError("Minimal bin width too large for the number of bins")
|
||||
if min_bin_height * num_bins > 1.0:
|
||||
raise ValueError("Minimal bin height too large for the number of bins")
|
||||
|
||||
widths = F.softmax(unnormalized_widths, dim=-1)
|
||||
widths = min_bin_width + (1 - min_bin_width * num_bins) * widths
|
||||
cumwidths = torch.cumsum(widths, dim=-1)
|
||||
cumwidths = F.pad(cumwidths, pad=(1, 0), mode="constant", value=0.0)
|
||||
cumwidths = (right - left) * cumwidths + left
|
||||
cumwidths[..., 0] = left
|
||||
cumwidths[..., -1] = right
|
||||
widths = cumwidths[..., 1:] - cumwidths[..., :-1]
|
||||
|
||||
derivatives = min_derivative + F.softplus(unnormalized_derivatives)
|
||||
|
||||
heights = F.softmax(unnormalized_heights, dim=-1)
|
||||
heights = min_bin_height + (1 - min_bin_height * num_bins) * heights
|
||||
cumheights = torch.cumsum(heights, dim=-1)
|
||||
cumheights = F.pad(cumheights, pad=(1, 0), mode="constant", value=0.0)
|
||||
cumheights = (top - bottom) * cumheights + bottom
|
||||
cumheights[..., 0] = bottom
|
||||
cumheights[..., -1] = top
|
||||
heights = cumheights[..., 1:] - cumheights[..., :-1]
|
||||
|
||||
if inverse:
|
||||
bin_idx = searchsorted(cumheights, inputs)[..., None]
|
||||
else:
|
||||
bin_idx = searchsorted(cumwidths, inputs)[..., None]
|
||||
|
||||
input_cumwidths = cumwidths.gather(-1, bin_idx)[..., 0]
|
||||
input_bin_widths = widths.gather(-1, bin_idx)[..., 0]
|
||||
|
||||
input_cumheights = cumheights.gather(-1, bin_idx)[..., 0]
|
||||
delta = heights / widths
|
||||
input_delta = delta.gather(-1, bin_idx)[..., 0]
|
||||
|
||||
input_derivatives = derivatives.gather(-1, bin_idx)[..., 0]
|
||||
input_derivatives_plus_one = derivatives[..., 1:].gather(-1, bin_idx)[..., 0]
|
||||
|
||||
input_heights = heights.gather(-1, bin_idx)[..., 0]
|
||||
|
||||
if inverse:
|
||||
a = (inputs - input_cumheights) * (
|
||||
input_derivatives + input_derivatives_plus_one - 2 * input_delta
|
||||
) + input_heights * (input_delta - input_derivatives)
|
||||
b = input_heights * input_derivatives - (inputs - input_cumheights) * (
|
||||
input_derivatives + input_derivatives_plus_one - 2 * input_delta
|
||||
)
|
||||
c = -input_delta * (inputs - input_cumheights)
|
||||
|
||||
discriminant = b.pow(2) - 4 * a * c
|
||||
assert (discriminant >= 0).all()
|
||||
|
||||
root = (2 * c) / (-b - torch.sqrt(discriminant))
|
||||
outputs = root * input_bin_widths + input_cumwidths
|
||||
|
||||
theta_one_minus_theta = root * (1 - root)
|
||||
denominator = input_delta + (
|
||||
(input_derivatives + input_derivatives_plus_one - 2 * input_delta)
|
||||
* theta_one_minus_theta
|
||||
)
|
||||
derivative_numerator = input_delta.pow(2) * (
|
||||
input_derivatives_plus_one * root.pow(2)
|
||||
+ 2 * input_delta * theta_one_minus_theta
|
||||
+ input_derivatives * (1 - root).pow(2)
|
||||
)
|
||||
logabsdet = torch.log(derivative_numerator) - 2 * torch.log(denominator)
|
||||
|
||||
return outputs, -logabsdet
|
||||
else:
|
||||
theta = (inputs - input_cumwidths) / input_bin_widths
|
||||
theta_one_minus_theta = theta * (1 - theta)
|
||||
|
||||
numerator = input_heights * (
|
||||
input_delta * theta.pow(2) + input_derivatives * theta_one_minus_theta
|
||||
)
|
||||
denominator = input_delta + (
|
||||
(input_derivatives + input_derivatives_plus_one - 2 * input_delta)
|
||||
* theta_one_minus_theta
|
||||
)
|
||||
outputs = input_cumheights + numerator / denominator
|
||||
|
||||
derivative_numerator = input_delta.pow(2) * (
|
||||
input_derivatives_plus_one * theta.pow(2)
|
||||
+ 2 * input_delta * theta_one_minus_theta
|
||||
+ input_derivatives * (1 - theta).pow(2)
|
||||
)
|
||||
logabsdet = torch.log(derivative_numerator) - 2 * torch.log(denominator)
|
||||
|
||||
return outputs, logabsdet
|
||||
@@ -0,0 +1,19 @@
|
||||
from .symbols import *
|
||||
|
||||
|
||||
_symbol_to_id = {s: i for i, s in enumerate(symbols)}
|
||||
|
||||
|
||||
def phonemes_to_ids(cleaned_text, tones, language, symbol_to_id=None):
|
||||
"""Converts a list of phoneme symbols to a sequence of integer IDs."""
|
||||
symbol_to_id_map = symbol_to_id if symbol_to_id else _symbol_to_id
|
||||
unk_id = symbol_to_id_map.get("UNK")
|
||||
if unk_id is None:
|
||||
phones = [symbol_to_id_map[symbol] for symbol in cleaned_text]
|
||||
else:
|
||||
phones = [symbol_to_id_map.get(symbol, unk_id) for symbol in cleaned_text]
|
||||
tone_start = language_tone_start_map[language]
|
||||
tones = [i + tone_start for i in tones]
|
||||
lang_id = language_id_map[language]
|
||||
lang_ids = [lang_id for _ in phones]
|
||||
return phones, tones, lang_ids
|
||||
+129530
File diff suppressed because it is too large
Load Diff
Binary file not shown.
@@ -0,0 +1,173 @@
|
||||
import pickle
|
||||
import os
|
||||
import re
|
||||
from g2p_en import G2p
|
||||
|
||||
from . import symbols
|
||||
|
||||
from .english_utils.abbreviations import expand_abbreviations
|
||||
from .english_utils.time_norm import expand_time_english
|
||||
from .english_utils.number_norm import normalize_numbers
|
||||
|
||||
|
||||
def distribute_phone(n_phone, n_word):
|
||||
phones_per_word = [0] * n_word
|
||||
for task in range(n_phone):
|
||||
min_tasks = min(phones_per_word)
|
||||
min_indices = [
|
||||
i for i, x in enumerate(phones_per_word) if x == min_tasks
|
||||
]
|
||||
chosen_index = min_indices[len(min_indices) // 2]
|
||||
phones_per_word[chosen_index] += 1
|
||||
return phones_per_word
|
||||
|
||||
|
||||
from transformers import AutoTokenizer
|
||||
|
||||
current_file_path = os.path.dirname(__file__)
|
||||
CMU_DICT_PATH = os.path.join(current_file_path, "cmudict.rep")
|
||||
CACHE_PATH = os.path.join(current_file_path, "cmudict_cache.pickle")
|
||||
_g2p = G2p()
|
||||
|
||||
arpa = {
|
||||
"AH0", "S", "AH1", "EY2", "AE2", "EH0", "OW2", "UH0", "NG", "B",
|
||||
"G", "AY0", "M", "AA0", "F", "AO0", "ER2", "UH1", "IY1", "AH2",
|
||||
"DH", "IY0", "EY1", "IH0", "K", "N", "W", "IY2", "T", "AA1",
|
||||
"ER1", "EH2", "OY0", "UH2", "UW1", "Z", "AW2", "AW1", "V", "UW2",
|
||||
"AA2", "ER", "AW0", "UW0", "R", "OW1", "EH1", "ZH", "AE0", "IH2",
|
||||
"IH", "Y", "JH", "P", "AY1", "EY0", "OY2", "TH", "HH", "D",
|
||||
"ER0", "CH", "AO1", "AE1", "AO2", "OY1", "AY2", "IH1", "OW0", "L", "SH",
|
||||
}
|
||||
|
||||
|
||||
def map_phoneme(ph):
|
||||
rep_map = {
|
||||
"：": ",", "；": ",", "，": ",", "。": ".", "！": "!",
|
||||
"？": "?", "\n": ".", "·": ",", "、": ",", "...": "…", "v": "V",
|
||||
}
|
||||
if ph in rep_map.keys():
|
||||
ph = rep_map[ph]
|
||||
if ph in symbols:
|
||||
return ph
|
||||
if ph not in symbols:
|
||||
ph = "UNK"
|
||||
return ph
|
||||
|
||||
|
||||
def read_dict():
|
||||
g2p_dict = {}
|
||||
start_line = 49
|
||||
with open(CMU_DICT_PATH) as f:
|
||||
line = f.readline()
|
||||
line_index = 1
|
||||
while line:
|
||||
if line_index >= start_line:
|
||||
line = line.strip()
|
||||
word_split = line.split(" ")
|
||||
word = word_split[0]
|
||||
|
||||
syllable_split = word_split[1].split(" - ")
|
||||
g2p_dict[word] = []
|
||||
for syllable in syllable_split:
|
||||
phone_split = syllable.split(" ")
|
||||
g2p_dict[word].append(phone_split)
|
||||
|
||||
line_index = line_index + 1
|
||||
line = f.readline()
|
||||
|
||||
return g2p_dict
|
||||
|
||||
|
||||
def cache_dict(g2p_dict, file_path):
|
||||
with open(file_path, "wb") as pickle_file:
|
||||
pickle.dump(g2p_dict, pickle_file)
|
||||
|
||||
|
||||
def get_dict():
|
||||
if os.path.exists(CACHE_PATH):
|
||||
with open(CACHE_PATH, "rb") as pickle_file:
|
||||
g2p_dict = pickle.load(pickle_file)
|
||||
else:
|
||||
g2p_dict = read_dict()
|
||||
cache_dict(g2p_dict, CACHE_PATH)
|
||||
|
||||
return g2p_dict
|
||||
|
||||
|
||||
eng_dict = get_dict()
|
||||
|
||||
|
||||
def parse_phoneme(phn):
|
||||
tone = 0
|
||||
if re.search(r"\d$", phn):
|
||||
tone = int(phn[-1]) + 1
|
||||
phn = phn[:-1]
|
||||
return phn.lower(), tone
|
||||
|
||||
|
||||
def parse_syllables(syllables):
|
||||
tones = []
|
||||
phonemes = []
|
||||
for phn_list in syllables:
|
||||
for i in range(len(phn_list)):
|
||||
phn = phn_list[i]
|
||||
phn, tone = parse_phoneme(phn)
|
||||
phonemes.append(phn)
|
||||
tones.append(tone)
|
||||
return phonemes, tones
|
||||
|
||||
|
||||
def normalize_text(text):
|
||||
text = text.lower()
|
||||
text = expand_time_english(text)
|
||||
text = normalize_numbers(text)
|
||||
text = expand_abbreviations(text)
|
||||
return text
|
||||
|
||||
|
||||
model_id = 'bert-base-uncased'
|
||||
tokenizer = AutoTokenizer.from_pretrained(model_id)
|
||||
|
||||
|
||||
def grapheme_to_phoneme(text, pad_start_end=True, tokenized=None):
|
||||
if tokenized is None:
|
||||
tokenized = tokenizer.tokenize(text)
|
||||
ph_groups = []
|
||||
for t in tokenized:
|
||||
if not t.startswith("#"):
|
||||
ph_groups.append([t])
|
||||
else:
|
||||
ph_groups[-1].append(t.replace("#", ""))
|
||||
|
||||
phones = []
|
||||
tones = []
|
||||
word2ph = []
|
||||
for group in ph_groups:
|
||||
w = "".join(group)
|
||||
phone_len = 0
|
||||
word_len = len(group)
|
||||
if w.upper() in eng_dict:
|
||||
phns, tns = parse_syllables(eng_dict[w.upper()])
|
||||
phones += phns
|
||||
tones += tns
|
||||
phone_len += len(phns)
|
||||
else:
|
||||
phone_list = list(filter(lambda p: p != " ", _g2p(w)))
|
||||
for ph in phone_list:
|
||||
if ph in arpa:
|
||||
ph, tn = parse_phoneme(ph)
|
||||
phones.append(ph)
|
||||
tones.append(tn)
|
||||
else:
|
||||
phones.append(ph)
|
||||
tones.append(0)
|
||||
phone_len += 1
|
||||
aaa = distribute_phone(phone_len, word_len)
|
||||
word2ph += aaa
|
||||
phones = [map_phoneme(i) for i in phones]
|
||||
|
||||
if pad_start_end:
|
||||
phones = ["_"] + phones + ["_"]
|
||||
tones = [0] + tones + [0]
|
||||
word2ph = [1] + word2ph + [1]
|
||||
return phones, tones, word2ph
|
||||
@@ -0,0 +1,35 @@
|
||||
import re
|
||||
|
||||
# List of (regular expression, replacement) pairs for abbreviations in english:
|
||||
abbreviations_en = [
|
||||
(re.compile("\\b%s\\." % x[0], re.IGNORECASE), x[1])
|
||||
for x in [
|
||||
("mrs", "misess"),
|
||||
("mr", "mister"),
|
||||
("dr", "doctor"),
|
||||
("st", "saint"),
|
||||
("co", "company"),
|
||||
("jr", "junior"),
|
||||
("maj", "major"),
|
||||
("gen", "general"),
|
||||
("drs", "doctors"),
|
||||
("rev", "reverend"),
|
||||
("lt", "lieutenant"),
|
||||
("hon", "honorable"),
|
||||
("sgt", "sergeant"),
|
||||
("capt", "captain"),
|
||||
("esq", "esquire"),
|
||||
("ltd", "limited"),
|
||||
("col", "colonel"),
|
||||
("ft", "fort"),
|
||||
]
|
||||
]
|
||||
|
||||
def expand_abbreviations(text, lang="en"):
|
||||
if lang == "en":
|
||||
_abbreviations = abbreviations_en
|
||||
else:
|
||||
raise NotImplementedError()
|
||||
for regex, replacement in _abbreviations:
|
||||
text = re.sub(regex, replacement, text)
|
||||
return text
|
||||
@@ -0,0 +1,97 @@
|
||||
""" from https://github.com/keithito/tacotron """
|
||||
|
||||
import re
|
||||
from typing import Dict
|
||||
|
||||
import inflect
|
||||
|
||||
_inflect = inflect.engine()
|
||||
_comma_number_re = re.compile(r"([0-9][0-9\,]+[0-9])")
|
||||
_decimal_number_re = re.compile(r"([0-9]+\.[0-9]+)")
|
||||
_currency_re = re.compile(r"(£|\$|¥)([0-9\,\.]*[0-9]+)")
|
||||
_ordinal_re = re.compile(r"[0-9]+(st|nd|rd|th)")
|
||||
_number_re = re.compile(r"-?[0-9]+")
|
||||
|
||||
|
||||
def _remove_commas(m):
|
||||
return m.group(1).replace(",", "")
|
||||
|
||||
|
||||
def _expand_decimal_point(m):
|
||||
return m.group(1).replace(".", " point ")
|
||||
|
||||
|
||||
def __expand_currency(value: str, inflection: Dict[float, str]) -> str:
|
||||
parts = value.replace(",", "").split(".")
|
||||
if len(parts) > 2:
|
||||
return f"{value} {inflection[2]}" # Unexpected format
|
||||
text = []
|
||||
integer = int(parts[0]) if parts[0] else 0
|
||||
if integer > 0:
|
||||
integer_unit = inflection.get(integer, inflection[2])
|
||||
text.append(f"{integer} {integer_unit}")
|
||||
fraction = int(parts[1]) if len(parts) > 1 and parts[1] else 0
|
||||
if fraction > 0:
|
||||
fraction_unit = inflection.get(fraction / 100, inflection[0.02])
|
||||
text.append(f"{fraction} {fraction_unit}")
|
||||
if len(text) == 0:
|
||||
return f"zero {inflection[2]}"
|
||||
return " ".join(text)
|
||||
|
||||
|
||||
def _expand_currency(m: "re.Match") -> str:
|
||||
currencies = {
|
||||
"$": {
|
||||
0.01: "cent",
|
||||
0.02: "cents",
|
||||
1: "dollar",
|
||||
2: "dollars",
|
||||
},
|
||||
"€": {
|
||||
0.01: "cent",
|
||||
0.02: "cents",
|
||||
1: "euro",
|
||||
2: "euros",
|
||||
},
|
||||
"£": {
|
||||
0.01: "penny",
|
||||
0.02: "pence",
|
||||
1: "pound sterling",
|
||||
2: "pounds sterling",
|
||||
},
|
||||
"¥": {
|
||||
# TODO rin
|
||||
0.02: "sen",
|
||||
2: "yen",
|
||||
},
|
||||
}
|
||||
unit = m.group(1)
|
||||
currency = currencies[unit]
|
||||
value = m.group(2)
|
||||
return __expand_currency(value, currency)
|
||||
|
||||
|
||||
def _expand_ordinal(m):
|
||||
return _inflect.number_to_words(m.group(0))
|
||||
|
||||
|
||||
def _expand_number(m):
|
||||
num = int(m.group(0))
|
||||
if 1000 < num < 3000:
|
||||
if num == 2000:
|
||||
return "two thousand"
|
||||
if 2000 < num < 2010:
|
||||
return "two thousand " + _inflect.number_to_words(num % 100)
|
||||
if num % 100 == 0:
|
||||
return _inflect.number_to_words(num // 100) + " hundred"
|
||||
return _inflect.number_to_words(num, andword="", zero="oh", group=2).replace(", ", " ")
|
||||
return _inflect.number_to_words(num, andword="")
|
||||
|
||||
|
||||
def normalize_numbers(text):
|
||||
text = re.sub(_comma_number_re, _remove_commas, text)
|
||||
text = re.sub(_currency_re, _expand_currency, text)
|
||||
text = re.sub(_decimal_number_re, _expand_decimal_point, text)
|
||||
text = re.sub(_ordinal_re, _expand_ordinal, text)
|
||||
text = re.sub(_number_re, _expand_number, text)
|
||||
return text
|
||||
@@ -0,0 +1,47 @@
|
||||
import re
|
||||
|
||||
import inflect
|
||||
|
||||
_inflect = inflect.engine()
|
||||
|
||||
_time_re = re.compile(
|
||||
r"""\b
|
||||
((0?[0-9])|(1[0-1])|(1[2-9])|(2[0-3])) # hours
|
||||
:
|
||||
([0-5][0-9]) # minutes
|
||||
\s*(a\\.m\\.|am|pm|p\\.m\\.|a\\.m|p\\.m)? # am/pm
|
||||
\b""",
|
||||
re.IGNORECASE | re.X,
|
||||
)
|
||||
|
||||
|
||||
def _expand_num(n: int) -> str:
|
||||
return _inflect.number_to_words(n)
|
||||
|
||||
|
||||
def _expand_time_english(match: "re.Match") -> str:
|
||||
hour = int(match.group(1))
|
||||
past_noon = hour >= 12
|
||||
time = []
|
||||
if hour > 12:
|
||||
hour -= 12
|
||||
elif hour == 0:
|
||||
hour = 12
|
||||
past_noon = True
|
||||
time.append(_expand_num(hour))
|
||||
|
||||
minute = int(match.group(6))
|
||||
if minute > 0:
|
||||
if minute < 10:
|
||||
time.append("oh")
|
||||
time.append(_expand_num(minute))
|
||||
am_pm = match.group(7)
|
||||
if am_pm is None:
|
||||
time.append("p m" if past_noon else "a m")
|
||||
else:
|
||||
time.extend(list(am_pm.replace(".", "")))
|
||||
return " ".join(time)
|
||||
|
||||
|
||||
def expand_time_english(text: str) -> str:
|
||||
return re.sub(_time_re, _expand_time_english, text)
|
||||
@@ -0,0 +1,293 @@
|
||||
# punctuation = ["!", "?", "…", ",", ".", "'", "-"]
|
||||
punctuation = ["!", "?", "…", ",", ".", "'", "-", "¿", "¡"]
|
||||
pu_symbols = punctuation + ["SP", "UNK"]
|
||||
pad = "_"
|
||||
|
||||
# chinese
|
||||
zh_symbols = [
|
||||
"E",
|
||||
"En",
|
||||
"a",
|
||||
"ai",
|
||||
"an",
|
||||
"ang",
|
||||
"ao",
|
||||
"b",
|
||||
"c",
|
||||
"ch",
|
||||
"d",
|
||||
"e",
|
||||
"ei",
|
||||
"en",
|
||||
"eng",
|
||||
"er",
|
||||
"f",
|
||||
"g",
|
||||
"h",
|
||||
"i",
|
||||
"i0",
|
||||
"ia",
|
||||
"ian",
|
||||
"iang",
|
||||
"iao",
|
||||
"ie",
|
||||
"in",
|
||||
"ing",
|
||||
"iong",
|
||||
"ir",
|
||||
"iu",
|
||||
"j",
|
||||
"k",
|
||||
"l",
|
||||
"m",
|
||||
"n",
|
||||
"o",
|
||||
"ong",
|
||||
"ou",
|
||||
"p",
|
||||
"q",
|
||||
"r",
|
||||
"s",
|
||||
"sh",
|
||||
"t",
|
||||
"u",
|
||||
"ua",
|
||||
"uai",
|
||||
"uan",
|
||||
"uang",
|
||||
"ui",
|
||||
"un",
|
||||
"uo",
|
||||
"v",
|
||||
"van",
|
||||
"ve",
|
||||
"vn",
|
||||
"w",
|
||||
"x",
|
||||
"y",
|
||||
"z",
|
||||
"zh",
|
||||
"AA",
|
||||
"EE",
|
||||
"OO",
|
||||
]
|
||||
num_zh_tones = 6
|
||||
|
||||
# japanese
|
||||
ja_symbols = [
|
||||
"N",
|
||||
"a",
|
||||
"a:",
|
||||
"b",
|
||||
"by",
|
||||
"ch",
|
||||
"d",
|
||||
"dy",
|
||||
"e",
|
||||
"e:",
|
||||
"f",
|
||||
"g",
|
||||
"gy",
|
||||
"h",
|
||||
"hy",
|
||||
"i",
|
||||
"i:",
|
||||
"j",
|
||||
"k",
|
||||
"ky",
|
||||
"m",
|
||||
"my",
|
||||
"n",
|
||||
"ny",
|
||||
"o",
|
||||
"o:",
|
||||
"p",
|
||||
"py",
|
||||
"q",
|
||||
"r",
|
||||
"ry",
|
||||
"s",
|
||||
"sh",
|
||||
"t",
|
||||
"ts",
|
||||
"ty",
|
||||
"u",
|
||||
"u:",
|
||||
"w",
|
||||
"y",
|
||||
"z",
|
||||
"zy",
|
||||
]
|
||||
num_ja_tones = 1
|
||||
|
||||
# English
|
||||
en_symbols = [
|
||||
"aa",
|
||||
"ae",
|
||||
"ah",
|
||||
"ao",
|
||||
"aw",
|
||||
"ay",
|
||||
"b",
|
||||
"ch",
|
||||
"d",
|
||||
"dh",
|
||||
"eh",
|
||||
"er",
|
||||
"ey",
|
||||
"f",
|
||||
"g",
|
||||
"hh",
|
||||
"ih",
|
||||
"iy",
|
||||
"jh",
|
||||
"k",
|
||||
"l",
|
||||
"m",
|
||||
"n",
|
||||
"ng",
|
||||
"ow",
|
||||
"oy",
|
||||
"p",
|
||||
"r",
|
||||
"s",
|
||||
"sh",
|
||||
"t",
|
||||
"th",
|
||||
"uh",
|
||||
"uw",
|
||||
"V",
|
||||
"w",
|
||||
"y",
|
||||
"z",
|
||||
"zh",
|
||||
]
|
||||
num_en_tones = 4
|
||||
|
||||
# Korean
|
||||
kr_symbols = ['ᄌ', 'ᅥ', 'ᆫ', 'ᅦ', 'ᄋ', 'ᅵ', 'ᄅ', 'ᅴ', 'ᄀ', 'ᅡ', 'ᄎ', 'ᅪ', 'ᄑ', 'ᅩ', 'ᄐ', 'ᄃ', 'ᅢ', 'ᅮ', 'ᆼ', 'ᅳ', 'ᄒ', 'ᄆ', 'ᆯ', 'ᆷ', 'ᄂ', 'ᄇ', 'ᄉ', 'ᆮ', 'ᄁ', 'ᅬ', 'ᅣ', 'ᄄ', 'ᆨ', 'ᄍ', 'ᅧ', 'ᄏ', 'ᆸ', 'ᅭ', '(', 'ᄊ', ')', 'ᅲ', 'ᅨ', 'ᄈ', 'ᅱ', 'ᅯ', 'ᅫ', 'ᅰ', 'ᅤ', '~', '\\', '[', ']', '/', '^', ':', 'ㄸ', '*']
|
||||
num_kr_tones = 1
|
||||
|
||||
# Spanish
|
||||
es_symbols = [
|
||||
"N",
|
||||
"Q",
|
||||
"a",
|
||||
"b",
|
||||
"d",
|
||||
"e",
|
||||
"f",
|
||||
"g",
|
||||
"h",
|
||||
"i",
|
||||
"j",
|
||||
"k",
|
||||
"l",
|
||||
"m",
|
||||
"n",
|
||||
"o",
|
||||
"p",
|
||||
"s",
|
||||
"t",
|
||||
"u",
|
||||
"v",
|
||||
"w",
|
||||
"x",
|
||||
"y",
|
||||
"z",
|
||||
"ɑ",
|
||||
"æ",
|
||||
"ʃ",
|
||||
"ʑ",
|
||||
"ç",
|
||||
"ɯ",
|
||||
"ɪ",
|
||||
"ɔ",
|
||||
"ɛ",
|
||||
"ɹ",
|
||||
"ð",
|
||||
"ə",
|
||||
"ɫ",
|
||||
"ɥ",
|
||||
"ɸ",
|
||||
"ʊ",
|
||||
"ɾ",
|
||||
"ʒ",
|
||||
"θ",
|
||||
"β",
|
||||
"ŋ",
|
||||
"ɦ",
|
||||
"ɡ",
|
||||
"r",
|
||||
"ɲ",
|
||||
"ʝ",
|
||||
"ɣ",
|
||||
"ʎ",
|
||||
"ˈ",
|
||||
"ˌ",
|
||||
"ː"
|
||||
]
|
||||
num_es_tones = 1
|
||||
|
||||
# French
|
||||
fr_symbols = [
|
||||
"\u0303",
|
||||
"œ",
|
||||
"ø",
|
||||
"ʁ",
|
||||
"ɒ",
|
||||
"ʌ",
|
||||
"ɜ",
|
||||
"ɐ"
|
||||
]
|
||||
num_fr_tones = 1
|
||||
|
||||
# German
|
||||
de_symbols = [
|
||||
"ʏ",
|
||||
"̩"
|
||||
]
|
||||
num_de_tones = 1
|
||||
|
||||
# Russian
|
||||
ru_symbols = [
|
||||
"ɭ",
|
||||
"ʲ",
|
||||
"ɕ",
|
||||
"\"",
|
||||
"ɵ",
|
||||
"^",
|
||||
"ɬ"
|
||||
]
|
||||
num_ru_tones = 1
|
||||
|
||||
# combine all symbols
|
||||
normal_symbols = sorted(set(zh_symbols + ja_symbols + en_symbols + kr_symbols + es_symbols + fr_symbols + de_symbols + ru_symbols))
|
||||
symbols = [pad] + normal_symbols + pu_symbols
|
||||
sil_phonemes_ids = [symbols.index(i) for i in pu_symbols]
|
||||
|
||||
# combine all tones
|
||||
num_tones = num_zh_tones + num_ja_tones + num_en_tones + num_kr_tones + num_es_tones + num_fr_tones + num_de_tones + num_ru_tones
|
||||
|
||||
# language maps
|
||||
language_id_map = {"ZH": 0, "JP": 1, "EN": 2, "ZH_MIX_EN": 3, 'KR': 4, 'ES': 5, 'SP': 5, 'FR': 6, 'DE': 7, 'RU': 8, 'VI': 9}
|
||||
num_languages = 10
|
||||
|
||||
language_tone_start_map = {
|
||||
"ZH": 0,
|
||||
"ZH_MIX_EN": 0,
|
||||
"JP": num_zh_tones,
|
||||
"EN": num_zh_tones + num_ja_tones,
|
||||
'KR': num_zh_tones + num_ja_tones + num_en_tones,
|
||||
"ES": num_zh_tones + num_ja_tones + num_en_tones + num_kr_tones,
|
||||
"SP": num_zh_tones + num_ja_tones + num_en_tones + num_kr_tones,
|
||||
"FR": num_zh_tones + num_ja_tones + num_en_tones + num_kr_tones + num_es_tones,
|
||||
"DE": num_zh_tones + num_ja_tones + num_en_tones + num_kr_tones + num_es_tones + num_fr_tones,
|
||||
"RU": num_zh_tones + num_ja_tones + num_en_tones + num_kr_tones + num_es_tones + num_fr_tones + num_de_tones,
|
||||
"VI": num_zh_tones + num_ja_tones + num_en_tones + num_kr_tones + num_es_tones + num_fr_tones + num_de_tones + num_ru_tones,
|
||||
}
|
||||
|
||||
if __name__ == "__main__":
|
||||
a = set(zh_symbols)
|
||||
b = set(en_symbols)
|
||||
print(sorted(a & b))
|
||||
@@ -0,0 +1,5 @@
|
||||
from .config import (
|
||||
SAMPLING_RATE, FILTER_LENGTH, HOP_LENGTH, SEGMENT_FRAMES,
|
||||
ADD_BLANK, SPEC_CHANNELS, N_SPEAKERS, SPK2ID,
|
||||
MODEL_PARAMS, NUM_LANGUAGES, NUM_TONES,
|
||||
)
|
||||
@@ -0,0 +1,42 @@
|
||||
# Audio
|
||||
SAMPLING_RATE = 44100
|
||||
FILTER_LENGTH = 2048
|
||||
HOP_LENGTH = 512
|
||||
SEGMENT_FRAMES = 32
|
||||
ADD_BLANK = True
|
||||
SPEC_CHANNELS = FILTER_LENGTH // 2 + 1 # 1025
|
||||
N_MEL_CHANNELS = 128 # updated in new checkpoint
|
||||
|
||||
# Speakers
|
||||
N_SPEAKERS = 1
|
||||
SPK2ID = {"MALE": 0}
|
||||
|
||||
# Model — matches config.json for G_150000.pth (lighter version)
|
||||
MODEL_PARAMS = dict(
|
||||
use_spk_conditioned_encoder=True,
|
||||
use_noise_scaled_mas=True,
|
||||
inter_channels=32,
|
||||
hidden_channels=32,
|
||||
filter_channels=128,
|
||||
n_heads=2,
|
||||
n_layers=3,
|
||||
n_layers_trans_flow=3,
|
||||
kernel_size=3,
|
||||
p_dropout=0.1,
|
||||
resblock="1",
|
||||
resblock_kernel_sizes=[3, 7, 11],
|
||||
resblock_dilation_sizes=[[1, 3, 5], [1, 3, 5], [1, 3, 5]],
|
||||
upsample_rates=[8, 8, 2, 2, 2],
|
||||
upsample_initial_channel=64,
|
||||
upsample_kernel_sizes=[16, 16, 8, 2, 2],
|
||||
n_layers_q=3,
|
||||
use_spectral_norm=False,
|
||||
gin_channels=128,
|
||||
use_sdp=True,
|
||||
mas_noise_scale_initial=0.01,
|
||||
noise_scale_delta=2e-06,
|
||||
)
|
||||
|
||||
# Language / Tone
|
||||
NUM_LANGUAGES = 1
|
||||
NUM_TONES = 6
|
||||
Reference in New Issue
Block a user