| |
| """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) |
|
|
|
|
| |
| 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): |
| |
| |
| 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) |
|
|
|
|
|
|
