File size: 6,941 Bytes
80be62c | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 | #!/usr/bin/env python3
"""The Swar network, built from Kokoro's own layer classes (swar/kokoro, Apache-2.0).
from swar.model import build
m = build("D", config) # dict of nn.Modules, widths from config.json
Stage C trains text_encoder + decoder + style_encoder (acoustic); stage D adds
plbert + bert_encoder + predictor + predictor_encoder (prosodic style). At
inference the two 128-d styles concatenate into Kokoro's 256-d ref_s layout:
[:128] drives the decoder (timbre), [128:] drives the prosody predictor.
Two deliberate departures from the Kokoro code, both because we run at 16 kHz:
* Kokoro's Generator hard-codes SourceModuleHnNSF(sampling_rate=24000). Its sine
source turns F0 into a waveform at that rate, so at 16 kHz every voice would
come out 1.5x too high. The source is rebuilt at our rate.
* Kokoro's Decoder hard-codes its widths (1024 / 512 / 64); they come from
config.json here (512 / 256 / 64).
The style encoder is StyleTTS2's (Kokoro ships none -- it ships voice packs).
"""
from __future__ import annotations
import json
import math
from pathlib import Path
import torch
import torch.nn.functional as F
from torch import nn
from torch.nn.utils import spectral_norm, weight_norm
from .kokoro.istftnet import AdainResBlk1d, Generator, SourceModuleHnNSF
from .kokoro.modules import ProsodyPredictor, TextEncoder
SR = 16000
def conf(path: str | Path | None = None) -> dict:
p = Path(path) if path else Path(__file__).resolve().parents[1] / "config.json"
return json.loads(p.read_text())
class Decoder(nn.Module):
"""Kokoro's istftnet.Decoder: aligned text features + F0 + energy + style -> waveform."""
def __init__(self, dim_in, style_dim, hid, out, res, gen: dict, sr: int = SR):
super().__init__()
self.encode = AdainResBlk1d(dim_in + 2, hid, style_dim)
self.decode = nn.ModuleList([
AdainResBlk1d(hid + 2 + res, hid, style_dim),
AdainResBlk1d(hid + 2 + res, hid, style_dim),
AdainResBlk1d(hid + 2 + res, hid, style_dim),
AdainResBlk1d(hid + 2 + res, out, style_dim, upsample=True),
])
self.F0_conv = weight_norm(nn.Conv1d(1, 1, kernel_size=3, stride=2, groups=1, padding=1))
self.N_conv = weight_norm(nn.Conv1d(1, 1, kernel_size=3, stride=2, groups=1, padding=1))
self.asr_res = nn.Sequential(weight_norm(nn.Conv1d(dim_in, res, kernel_size=1)))
self.generator = Generator(style_dim, **gen)
self.generator.m_source = SourceModuleHnNSF(
sampling_rate=sr,
upsample_scale=math.prod(gen["upsample_rates"]) * gen["gen_istft_hop_size"],
harmonic_num=8, voiced_threshod=10)
def forward(self, asr, F0_curve, N, s):
"""asr [B, C, T] (40 Hz), F0/N [B, 2T] (80 Hz), s [B, style] -> [B, 1, 400 T]."""
F0 = self.F0_conv(F0_curve.unsqueeze(1))
N = self.N_conv(N.unsqueeze(1))
x = torch.cat([asr, F0, N], axis=1)
x = self.encode(x, s)
asr_res = self.asr_res(asr)
res = True
for block in self.decode:
if res:
x = torch.cat([x, asr_res, F0, N], axis=1)
x = block(x, s)
if block.upsample_type != "none":
res = False
return self.generator(x, s, F0_curve)
# ---- StyleTTS2 style encoder (Modules from yl4579/StyleTTS2 models.py, MIT) --------
class LearnedDownSample(nn.Module):
def __init__(self, dim_in):
super().__init__()
self.conv = spectral_norm(nn.Conv2d(dim_in, dim_in, kernel_size=(3, 3), stride=(2, 2),
groups=dim_in, padding=1))
def forward(self, x):
return self.conv(x)
class ResBlk(nn.Module):
def __init__(self, dim_in, dim_out):
super().__init__()
self.actv = nn.LeakyReLU(0.2)
self.learned_sc = dim_in != dim_out
self.conv1 = spectral_norm(nn.Conv2d(dim_in, dim_in, 3, 1, 1))
self.conv2 = spectral_norm(nn.Conv2d(dim_in, dim_out, 3, 1, 1))
self.down_res = LearnedDownSample(dim_in)
if self.learned_sc:
self.conv1x1 = spectral_norm(nn.Conv2d(dim_in, dim_out, 1, 1, 0, bias=False))
def _shortcut(self, x):
if self.learned_sc:
x = self.conv1x1(x)
return F.avg_pool2d(x, 2)
def _residual(self, x):
x = self.conv1(self.actv(x))
x = self.down_res(x)
return self.conv2(self.actv(x))
def forward(self, x):
return (self._shortcut(x) + self._residual(x)) / math.sqrt(2)
class StyleEncoder(nn.Module):
"""mel [B, 1, 80, T] -> [B, style_dim]. T >= 80 frames (1 s) recommended."""
def __init__(self, dim_in=32, style_dim=128, max_conv_dim=256):
super().__init__()
blocks = [spectral_norm(nn.Conv2d(1, dim_in, 3, 1, 1))]
d = dim_in
for _ in range(4):
o = min(d * 2, max_conv_dim)
blocks.append(ResBlk(d, o))
d = o
blocks += [nn.LeakyReLU(0.2), spectral_norm(nn.Conv2d(d, d, 5, 1, 0)),
nn.AdaptiveAvgPool2d(1), nn.LeakyReLU(0.2)]
self.shared = nn.Sequential(*blocks)
self.unshared = nn.Linear(d, style_dim)
def forward(self, x):
# 4 halvings: time must be a multiple of 16 or the pooled shortcut and the
# strided residual disagree by a frame (2 s of audio = 161 mel frames).
x = x[..., : (x.shape[-1] // 16) * 16]
return self.unshared(self.shared(x).flatten(1))
def build(stage: str = "C", c: dict | None = None) -> nn.ModuleDict:
c = c or conf()
se = c["style_encoder"]
m = {
"text_encoder": TextEncoder(c["hidden_dim"], c["text_encoder_kernel_size"], c["n_layer"],
c["n_token"]),
"decoder": Decoder(c["hidden_dim"], c["style_dim"], c["decoder"]["hid"], c["decoder"]["out"],
c["decoder"]["res"], c["istftnet"], sr=c["sample_rate"]),
"style_encoder": StyleEncoder(se["dim_in"], c["style_dim"], se["max_conv_dim"]),
}
if stage != "C":
from transformers import AlbertConfig, AlbertModel
m["plbert"] = AlbertModel(AlbertConfig(vocab_size=c["n_token"], hidden_act="gelu", **c["plbert"]),
add_pooling_layer=False)
m["bert_encoder"] = nn.Linear(c["plbert"]["hidden_size"], c["hidden_dim"])
m["predictor"] = ProsodyPredictor(c["style_dim"], c["hidden_dim"], c["n_layer"], c["max_dur"],
c["dropout"])
m["predictor_encoder"] = StyleEncoder(se["dim_in"], c["style_dim"], se["max_conv_dim"])
return nn.ModuleDict(m)
def energy(mel_norm: torch.Tensor) -> torch.Tensor:
"""StyleTTS2's N target: log of the L2 norm of the LINEAR mel frame, [B, 80, T] -> [B, T]."""
from .mel import MEAN, STD
return torch.log(torch.exp(mel_norm * STD + MEAN).norm(dim=1) + 1e-5)
|