#!/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)