https://github.com/audeering/shift

tts.py

@@ -0,0 +1,847 @@
+import torch
+import nltk
+nltk.download('punkt', download_dir='./')      # COMMENT IF DOWNLOADED
+nltk.download('punkt_tab', download_dir='./')  # COMMENT IF DOWNLOADED
+nltk.data.path.append('.')
+import librosa
+import audiofile
+import torch.nn.functional as F
+import math
+import numpy as np
+import torch.nn as nn
+import string
+import textwrap
+import phonemizer
+from espeak_util import set_espeak_library
+from transformers import AlbertConfig, AlbertModel
+from huggingface_hub import hf_hub_download
+from nltk.tokenize import word_tokenize
+from torch.nn import Conv1d, ConvTranspose1d
+from torch.nn.utils.parametrizations import weight_norm
+from torch.nn.utils import spectral_norm
+
+_pad = "$"
+_punctuation = ';:,.!?¡¿—…"«»“” '
+_letters = 'ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz'
+_letters_ipa = "ɑɐɒæɓʙβɔɕçɗɖðʤəɘɚɛɜɝɞɟʄɡɠɢʛɦɧħɥʜɨɪʝɭɬɫɮʟɱɯɰŋɳɲɴøɵɸθœɶʘɹɺɾɻʀʁɽʂʃʈʧʉʊʋⱱʌɣɤʍχʎʏʑʐʒʔʡʕʢǀǁǂǃˈˌːˑʼʴʰʱʲʷˠˤ˞↓↑→↗↘'̩'ᵻ"
+MAX_PHONEMES = 424   # For OOM is the max length of single (non-split) sentence for StyleTTS2 inference
+
+symbols = [_pad] + list(_punctuation) + list(_letters) + list(_letters_ipa)
+
+dicts = {}
+for i in range(len((symbols))):
+    dicts[symbols[i]] = i
+
+
+class TextCleaner:
+    def __init__(self, dummy=None):
+        self.word_index_dictionary = dicts
+        print(len(dicts))
+
+    def __call__(self, text):
+        indexes = []
+        for char in text:
+            try:
+                indexes.append(self.word_index_dictionary[char])
+            except KeyError:
+                # `=NONVOCAL    == \x00\x01\x02\x03\x04\x05\x06\x07\x08\t\n\x0b\x0c\r\x0e\x0f\x10\x11\x12\x13\x14\x15\x16\x17\x18\x19\x1a\x1b\x1c\x1d\x1e\x1f !"#$%&'()*+,-./0123456789:;<=>?@ABCDEFGHIJKLMNOPQRSTUVWXYZ[\]^_`abcdefghijklmnopqrstuvwxyz{|}~\x7f
+                # print(f'NonVOCAL {char}', end='\r')
+                pass
+        return indexes
+
+set_espeak_library()
+
+textclenaer = TextCleaner()
+
+global_phonemizer = phonemizer.backend.EspeakBackend(language="en-us", preserve_punctuation=True, with_stress=True)
+
+def _del_prefix(d):
+    # del ".module"
+    out = {}
+    for k, v in d.items():
+        out[k[7:]] = v
+    return out
+
+
+
+
+class StyleTTS2(nn.Module):
+
+    def __init__(self):
+        super().__init__()
+        albert_base_configuration = AlbertConfig(vocab_size=178,
+                                                 hidden_size=768,
+                                                 num_attention_heads=12,
+                                                 intermediate_size=2048,
+                                                 max_position_embeddings=512,
+                                                 num_hidden_layers=12,
+                                                 dropout=0.1)
+        self.bert = AlbertModel(albert_base_configuration)
+        state_dict = torch.load(hf_hub_download(repo_id='dkounadis/artificial-styletts2',
+                                                filename='Utils/PLBERT/step_1000000.pth'),
+                                map_location='cpu')['net']
+        new_state_dict = {}
+        for k, v in state_dict.items():
+            name = k[7:] # remove `module.`
+            if name.startswith('encoder.'):
+                name = name[8:] # remove `encoder.`
+                new_state_dict[name] = v
+        del new_state_dict["embeddings.position_ids"]
+        self.bert.load_state_dict(new_state_dict, strict=True)
+        self.decoder = Decoder(dim_in=512,
+                        style_dim=128,
+                        dim_out=80,  # n_mels
+                        resblock_kernel_sizes=[3, 7, 11],
+                        upsample_rates=[10, 5, 3, 2],
+                        upsample_initial_channel=512,
+                        resblock_dilation_sizes=[[1, 3, 5], [1, 3, 5], [1, 3, 5]],
+                        upsample_kernel_sizes=[20, 10, 6, 4])
+        self.text_encoder = TextEncoder(channels=512,
+                                kernel_size=5,
+                                depth=3,  # args['model_params']['n_layer'],
+                                n_symbols=178,  # args['model_params']['n_token']
+                                )
+        self.predictor = ProsodyPredictor(style_dim=128,
+                                    d_hid=512,
+                                    nlayers=3,  # OFFICIAL config.nlayers=5;
+                                    max_dur=50)
+        self.style_encoder = StyleEncoder()
+        self.predictor_encoder = StyleEncoder()
+        self.bert_encoder = torch.nn.Linear(self.bert.config.hidden_size, 512)
+        self.mel_spec = MelSpec()
+        params = torch.load(hf_hub_download(repo_id='yl4579/StyleTTS2-LibriTTS',
+                                            filename='Models/LibriTTS/epochs_2nd_00020.pth'),
+                            map_location='cpu')['net']
+        self.bert.load_state_dict(_del_prefix(params['bert']), strict=True)
+        self.bert_encoder.load_state_dict(_del_prefix(params['bert_encoder']), strict=True)
+        self.predictor.load_state_dict(_del_prefix(params['predictor']), strict=True)
+        self.decoder.load_state_dict(_del_prefix(params['decoder']), strict=True)
+        self.text_encoder.load_state_dict(_del_prefix(params['text_encoder']), strict=True)
+        self.predictor_encoder.load_state_dict(_del_prefix(params['predictor_encoder']), strict=True)
+        self.style_encoder.load_state_dict(_del_prefix(params['style_encoder']), strict=True)
+        
+        # FOR LSTM
+        for n, p in self.named_parameters():
+                p.requires_grad = False
+        self.eval()
+
+                
+    def device(self):
+        return self.style_encoder.unshared.weight.device
+
+    def compute_style(self, wav_file=None):
+        
+        x, sr = librosa.load(wav_file, sr=24000)
+        x, _ = librosa.effects.trim(x, top_db=30)
+        if sr != 24000:
+            x = librosa.resample(x, sr, 24000)
+        # LOGMEL - Has 16KHz default basisc - Called on 24KHz .wav
+        x = torch.from_numpy(x[None, :]).to(device=self.device(), 
+                                            dtype=torch.float)
+        mel_tensor = (torch.log(1e-5 + self.mel_spec(x)) + 4) / 4
+        #mel_tensor = preprocess(audio).to(device)
+        ref_s = self.style_encoder(mel_tensor)
+        ref_p = self.predictor_encoder(mel_tensor)  # [bs, 11, 1, 128]
+        s = torch.cat([ref_s, ref_p], dim=3)  # [bs, 11, 1, 256]
+        s = s[:, :, 0, :].transpose(1, 2)  # [1, 128, 11]
+        return s  # [1, 128, 11]
+        
+    def inference(self,
+                  text,
+                  ref_s=None):
+        '''text may become too long when phonemized'''
+
+        if isinstance(ref_s, str):
+            ref_s = self.compute_style(ref_s)
+        else:
+            pass # assume ref_s = precomputed style vector
+        
+        
+        # text = transliterate_number(text, lang='en').strip()
+        # as we are in english transliteration is already done by the text cleaner?
+        # somehow we have phonemes in text that try to be rephonemized
+        # The ds txt should be only ascii
+        
+        
+        if isinstance(text, str):
+            
+            _translator = str.maketrans('', '', string.punctuation)
+            
+            text = [sub_sent.translate(_translator) + '.' for sub_sent in textwrap.wrap(text, 74)]
+
+            # # text = nltk.sent_tokenize(text)
+            # # text = [i for sent in sentences for i in textwrap.wrap(sent, width=120)]
+            
+
+            # # text = textwrap.wrap(text, width=MAX_PHONEMES)  # phonemes thus sent_tokenize() can't split them in sentences
+            
+        
+        device = ref_s.device
+        total = []
+        for _t in text:
+            
+            _t = global_phonemizer.phonemize([_t])
+            _t = word_tokenize(_t[0])
+            _t = ' '.join(_t)
+            
+            tokens = textclenaer(_t)[:MAX_PHONEMES] + [4]  # textclenaer('.;?!') = [4,1,6,5] # append . punctuation to assure proper sound termination (pulse Issue)
+            
+            # After filter we should assure is terminating as a sentence
+            # print(len(_t), len(tokens), 'Msi')#, textclenaer('.;?!'))
+            # ================================= Delete Phonemes If len(phonemes) > len(text)  === OOM during training
+            tokens.insert(0, 0)
+            tokens = torch.LongTensor(tokens).to(device).unsqueeze(0)
+            with torch.no_grad():
+                hidden_states = self.text_encoder(tokens)
+                bert_dur = self.bert(tokens, attention_mask=torch.ones_like(tokens)
+                                     ).last_hidden_state
+                d_en = self.bert_encoder(bert_dur).transpose(-1, -2)
+                aln_trg, F0_pred, N_pred = self.predictor(d_en=d_en, s=ref_s[:, 128:, :])
+                asr = torch.bmm(aln_trg, hidden_states)
+                asr = asr.transpose(1, 2)
+                asr_new = torch.zeros_like(asr)
+                asr_new[:, :, 0] = asr[:, :, 0]
+                asr_new[:, :, 1:] = asr[:, :, 0:-1]
+                asr = asr_new
+                x = self.decoder(asr=asr,
+                            F0_curve=F0_pred,
+                            N=N_pred,
+                            s=ref_s[:, :128, :])  # different part of ref_s
+                # print(x.shape, 'TTS TTS TTS TTS')
+                if x.shape[2] < 100:
+                    x = torch.zeros(1, 1, 1000, device=self.device())  # silence if this sentence was empty
+                    
+            # NORMALIS / Crop Scratch at end (The endingscratch sound is not solved even with nltk.sentence split & punctuation)
+            x = x[..., 40:-4000]
+            # x /= x.abs().max() + 1e-7   # preserve as torch
+            # return x
+            if x.shape[2] == 0:
+                # nohing to vocode
+                x = torch.zeros(1, 1, 1000, device=self.device())
+            total.append(x)
+        
+        # --
+        total = 1.94 * torch.cat(total, 2)  # 1.94 * Perhaps exceeding -1,1 affects MIMI encode
+        total /= 1.02 * total.abs().max() + 1e-7
+        # --
+        return total
+
+
+
+
+def get_padding(kernel_size, dilation=1):
+    return int((kernel_size*dilation - dilation)/2)
+
+
+def _tile(x,
+          length=None):
+    x = x.repeat(1, 1, int(length / x.shape[2]) + 1)[:, :, :length]
+    return x
+
+
+class AdaIN1d(nn.Module):
+
+    # used by HiFiGan & ProsodyPredictor
+
+    def __init__(self, style_dim, num_features):
+        super().__init__()
+        self.norm = nn.InstanceNorm1d(num_features, affine=False)
+        self.fc = nn.Linear(style_dim, num_features*2)
+
+    def forward(self, x, s):
+
+        # x = torch.Size([1, 512, 248])     same as output
+        # s = torch.Size([1, 7, 1, 128])
+
+        s = self.fc(s.transpose(1, 2)).transpose(1, 2)
+
+        s = _tile(s, length=x.shape[2])
+
+        gamma, beta = torch.chunk(s, chunks=2, dim=1)
+        return (1+gamma) * self.norm(x) + beta
+
+
+class AdaINResBlock1(torch.nn.Module):
+    def __init__(self, channels, kernel_size=3, dilation=(1, 3, 5), style_dim=64):
+        super(AdaINResBlock1, self).__init__()
+        self.convs1 = nn.ModuleList([
+            weight_norm(Conv1d(channels, channels, kernel_size, 1, dilation=dilation[0],
+                               padding=get_padding(kernel_size, dilation[0]))),
+            weight_norm(Conv1d(channels, channels, kernel_size, 1, dilation=dilation[1],
+                               padding=get_padding(kernel_size, dilation[1]))),
+            weight_norm(Conv1d(channels, channels, kernel_size, 1, dilation=dilation[2],
+                               padding=get_padding(kernel_size, dilation[2])))
+        ])
+        # self.convs1.apply(init_weights)
+
+        self.convs2 = nn.ModuleList([
+            weight_norm(Conv1d(channels, channels, kernel_size, 1, dilation=1,
+                               padding=get_padding(kernel_size, 1))),
+            weight_norm(Conv1d(channels, channels, kernel_size, 1, dilation=1,
+                               padding=get_padding(kernel_size, 1))),
+            weight_norm(Conv1d(channels, channels, kernel_size, 1, dilation=1,
+                               padding=get_padding(kernel_size, 1)))
+        ])
+        # self.convs2.apply(init_weights)
+
+        self.adain1 = nn.ModuleList([
+            AdaIN1d(style_dim, channels),
+            AdaIN1d(style_dim, channels),
+            AdaIN1d(style_dim, channels),
+        ])
+
+        self.adain2 = nn.ModuleList([
+            AdaIN1d(style_dim, channels),
+            AdaIN1d(style_dim, channels),
+            AdaIN1d(style_dim, channels),
+        ])
+
+        self.alpha1 = nn.ParameterList(
+            [nn.Parameter(torch.ones(1, channels, 1)) for i in range(len(self.convs1))])
+        self.alpha2 = nn.ParameterList(
+            [nn.Parameter(torch.ones(1, channels, 1)) for i in range(len(self.convs2))])
+
+    def forward(self, x, s):
+        for c1, c2, n1, n2, a1, a2 in zip(self.convs1, self.convs2, self.adain1, self.adain2, self.alpha1, self.alpha2):
+            xt = n1(x, s)  # THIS IS ADAIN - EXPECTS conv1d dims
+            xt = xt + (1 / a1) * (torch.sin(a1 * xt) ** 2)  # Snake1D
+            xt = c1(xt)
+            xt = n2(xt, s)  # THIS IS ADAIN - EXPECTS conv1d dims
+            xt = xt + (1 / a2) * (torch.sin(a2 * xt) ** 2)  # Snake1D
+            xt = c2(xt)
+            x = xt + x
+        return x
+
+
+class SourceModuleHnNSF(torch.nn.Module):
+
+    def __init__(self):
+
+        super().__init__()
+        self.harmonic_num = 8
+        self.l_linear = torch.nn.Linear(self.harmonic_num + 1, 1)
+        self.upsample_scale = 300
+        
+
+    def forward(self, x):
+        # --
+        x = torch.multiply(x, torch.FloatTensor(
+            [[range(1, self.harmonic_num + 2)]]).to(x.device))  # [1, 145200, 9]
+        
+        # modulo of negative f0_values => -21 % 10 = 9 as -3*10 + 9 = 21 NOTICE THAT f0_values IS SIGNED
+        rad_values = x / 25647 #).clamp(0, 1)
+        # rad_values = torch.where(torch.logical_or(rad_values < 0, rad_values > 1), 0.5, rad_values)
+        rad_values = rad_values % 1  # % of neg values
+        rad_values = F.interpolate(rad_values.transpose(1, 2),
+                                                     scale_factor=1/self.upsample_scale,
+                                                     mode='linear').transpose(1, 2)
+
+        # 1.89 sounds also nice has woofer at punctuation
+        phase = torch.cumsum(rad_values, dim=1) * 1.84 * np.pi
+        phase = F.interpolate(phase.transpose(1, 2) * self.upsample_scale,
+                              scale_factor=self.upsample_scale, mode='linear').transpose(1, 2)
+        x = .009 * phase.sin()
+        # --
+        x = self.l_linear(x).tanh()
+        return x
+
+
+class Generator(torch.nn.Module):
+    def __init__(self,
+                 style_dim,
+                 resblock_kernel_sizes,
+                 upsample_rates,
+                 upsample_initial_channel,
+                 resblock_dilation_sizes,
+                 upsample_kernel_sizes):
+        super(Generator, self).__init__()
+        self.num_kernels = len(resblock_kernel_sizes)
+        self.num_upsamples = len(upsample_rates)
+        self.m_source = SourceModuleHnNSF()
+        self.f0_upsamp = torch.nn.Upsample(scale_factor=np.prod(upsample_rates))
+        self.noise_convs = nn.ModuleList()
+        self.ups = nn.ModuleList()
+        self.noise_res = nn.ModuleList()
+
+        for i, (u, k) in enumerate(zip(upsample_rates, upsample_kernel_sizes)):
+            c_cur = upsample_initial_channel // (2 ** (i + 1))
+
+            self.ups.append(weight_norm(ConvTranspose1d(upsample_initial_channel//(2**i),
+                                                        upsample_initial_channel//(
+                                                            2**(i+1)),
+                                                        k, u, padding=(u//2 + u % 2), output_padding=u % 2)))
+
+            if i + 1 < len(upsample_rates):
+                stride_f0 = np.prod(upsample_rates[i + 1:])
+                self.noise_convs.append(Conv1d(
+                    1, c_cur, kernel_size=stride_f0 * 2, stride=stride_f0, padding=(stride_f0+1) // 2))
+                self.noise_res.append(AdaINResBlock1(
+                    c_cur, 7, [1, 3, 5], style_dim))
+            else:
+                self.noise_convs.append(Conv1d(1, c_cur, kernel_size=1))
+                self.noise_res.append(AdaINResBlock1(
+                    c_cur, 11, [1, 3, 5], style_dim))
+
+        self.resblocks = nn.ModuleList()
+
+        self.alphas = nn.ParameterList()
+        self.alphas.append(nn.Parameter(
+            torch.ones(1, upsample_initial_channel, 1)))
+
+        for i in range(len(self.ups)):
+            ch = upsample_initial_channel//(2**(i+1))
+            self.alphas.append(nn.Parameter(torch.ones(1, ch, 1)))
+
+            for j, (k, d) in enumerate(zip(resblock_kernel_sizes, resblock_dilation_sizes)):
+                self.resblocks.append(AdaINResBlock1(ch, k, d, style_dim))
+
+        self.conv_post = weight_norm(Conv1d(ch, 1, 7, 1, padding=3))
+
+    def forward(self, x, s, f0):
+
+        # x.shape=torch.Size([1, 512, 484]) s.shape=torch.Size([1, 1, 1, 128]) f0.shape=torch.Size([1, 484]) GENERAT 249
+        f0 = self.f0_upsamp(f0).transpose(1, 2)
+
+        # x.shape=torch.Size([1, 512, 484]) s.shape=torch.Size([1, 1, 1, 128]) f0.shape=torch.Size([1, 145200, 1]) GENERAT 253
+
+        # [1, 145400, 1] f0 enters already upsampled to full wav 24kHz length
+        har_source = self.m_source(f0)
+
+        har_source = har_source.transpose(1, 2)
+
+        for i in range(self.num_upsamples):
+
+            x = x + (1 / self.alphas[i]) * (torch.sin(self.alphas[i] * x) ** 2)
+            x_source = self.noise_convs[i](har_source)
+            x_source = self.noise_res[i](x_source, s)
+
+            x = self.ups[i](x)
+
+            x = x + x_source
+
+            xs = None
+            for j in range(self.num_kernels):
+
+                if xs is None:
+                    xs = self.resblocks[i*self.num_kernels+j](x, s)
+                else:
+                    xs += self.resblocks[i*self.num_kernels+j](x, s)
+            x = xs / self.num_kernels
+        # x = x + (1 / self.alphas[i+1]) * (torch.sin(self.alphas[i+1] * x) ** 2)  # noisy
+        x = self.conv_post(x)
+        x = torch.tanh(x)
+
+        return x
+
+class AdainResBlk1d(nn.Module):
+
+    # also used in ProsodyPredictor()
+
+    def __init__(self, dim_in, dim_out, style_dim=64, actv=nn.LeakyReLU(0.2),
+                 upsample='none', dropout_p=0.0):
+        super().__init__()
+        self.actv = actv
+        self.upsample_type = upsample
+        self.upsample = UpSample1d(upsample)
+        self.learned_sc = dim_in != dim_out
+        self._build_weights(dim_in, dim_out, style_dim)
+        if upsample == 'none':
+            self.pool = nn.Identity()
+        else:
+            self.pool = weight_norm(nn.ConvTranspose1d(
+                dim_in, dim_in, kernel_size=3, stride=2, groups=dim_in, padding=1, output_padding=1))
+
+    def _build_weights(self, dim_in, dim_out, style_dim):
+        self.conv1 = weight_norm(nn.Conv1d(dim_in, dim_out, 3, 1, 1))
+        self.conv2 = weight_norm(nn.Conv1d(dim_out, dim_out, 3, 1, 1))
+        self.norm1 = AdaIN1d(style_dim, dim_in)
+        self.norm2 = AdaIN1d(style_dim, dim_out)
+        if self.learned_sc:
+            self.conv1x1 = weight_norm(
+                nn.Conv1d(dim_in, dim_out, 1, 1, 0, bias=False))
+
+    def _shortcut(self, x):
+        x = self.upsample(x)
+        if self.learned_sc:
+            x = self.conv1x1(x)
+        return x
+
+    def _residual(self, x, s):
+        x = self.norm1(x, s)
+        x = self.actv(x)
+        x = self.pool(x)
+        x = self.conv1(x)
+        x = self.norm2(x, s)
+        x = self.actv(x)
+        x = self.conv2(x)
+        return x
+
+    def forward(self, x, s):
+        out = self._residual(x, s)
+        out = (out + self._shortcut(x)) / math.sqrt(2)
+        return out
+
+
+class UpSample1d(nn.Module):
+    def __init__(self, layer_type):
+        super().__init__()
+        self.layer_type = layer_type
+
+    def forward(self, x):
+        if self.layer_type == 'none':
+            return x
+        else:
+            return F.interpolate(x, scale_factor=2, mode='nearest-exact')
+
+
+class Decoder(nn.Module):
+    def __init__(self, dim_in=512, F0_channel=512, style_dim=64, dim_out=80,
+                 resblock_kernel_sizes=[3, 7, 11],
+                 upsample_rates=[10, 5, 3, 2],
+                 upsample_initial_channel=512,
+                 resblock_dilation_sizes=[[1, 3, 5], [1, 3, 5], [1, 3, 5]],
+                 upsample_kernel_sizes=[20, 10, 6, 4]):
+        super().__init__()
+
+        self.decode = nn.ModuleList()
+
+        self.encode = AdainResBlk1d(dim_in + 2, 1024, style_dim)
+
+        self.decode.append(AdainResBlk1d(1024 + 2 + 64, 1024, style_dim))
+        self.decode.append(AdainResBlk1d(1024 + 2 + 64, 1024, style_dim))
+        self.decode.append(AdainResBlk1d(1024 + 2 + 64, 1024, style_dim))
+        self.decode.append(AdainResBlk1d(
+            1024 + 2 + 64, 512, style_dim, upsample=True))
+
+        self.F0_conv = weight_norm(
+            nn.Conv1d(1, 1, kernel_size=3, stride=2, groups=1, padding=1))  # smooth
+
+        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(512, 64, kernel_size=1)),
+        )
+
+        self.generator = Generator(style_dim, resblock_kernel_sizes, upsample_rates,
+                                   upsample_initial_channel, resblock_dilation_sizes, upsample_kernel_sizes)
+
+    def forward(self, asr=None, F0_curve=None, N=None, s=None):
+
+
+        F0 = self.F0_conv(F0_curve)
+        N = self.N_conv(N)
+
+
+        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
+
+        x = self.generator(x, s, F0_curve)
+        return x
+
+
+class MelSpec(torch.nn.Module):
+
+    def __init__(self,
+                 sample_rate=17402, # https://github.com/fakerybakery/styletts2-cli/blob/main/msinference.py = Default 16000. However 17400 vocalises better also "en_US/vctk_p274"
+                 n_fft=2048,
+                 win_length=1200,
+                 hop_length=300,
+                 n_mels=80
+                 ):
+        '''avoids dependency on torchaudio'''
+        super().__init__()
+        self.n_fft = n_fft
+        self.win_length = win_length if win_length is not None else n_fft
+        self.hop_length = hop_length if hop_length is not None else self.win_length // 2
+        # --
+        f_min = 0.0
+        f_max = float(sample_rate // 2)
+        all_freqs = torch.linspace(0, sample_rate // 2, n_fft//2+1)
+        m_min = 2595.0 * math.log10(1.0 + (f_min / 700.0))
+        m_max = 2595.0 * math.log10(1.0 + (f_max / 700.0))
+        m_pts = torch.linspace(m_min, m_max, n_mels + 2)
+        f_pts = 700.0 * (10 ** (m_pts / 2595.0) - 1.0)
+        f_diff = f_pts[1:] - f_pts[:-1]  # (n_mels + 1)
+        slopes = f_pts.unsqueeze(0) - all_freqs.unsqueeze(1)
+        zero = torch.zeros(1)
+        down_slopes = (-1.0 * slopes[:, :-2]) / f_diff[:-1]  # (n_freqs, n_mels)
+        up_slopes = slopes[:, 2:] / f_diff[1:]  # (n_freqs, n_mels)
+        fb = torch.max(zero, torch.min(down_slopes, up_slopes))
+        # --
+        self.register_buffer('fb', fb, persistent=False)
+        window = torch.hann_window(self.win_length)
+        self.register_buffer('window', window, persistent=False)
+
+    def forward(self, x):
+        spec_f = torch.stft(x,
+                            self.n_fft,
+                            self.hop_length,
+                            self.win_length,
+                            self.window,
+                            center=True,
+                            pad_mode="reflect",
+                            normalized=False,
+                            onesided=True,
+                            return_complex=True)  # [bs, 1025, 56]
+        mel_specgram = torch.matmul(spec_f.abs().pow(2).transpose(1, 2), self.fb).transpose(1, 2)
+        return mel_specgram[:, None, :, :]  # [bs, 1, 80, time]
+
+
+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)   # .07 also nice
+        self.downsample_res = LearnedDownSample(dim_in)
+        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))
+        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)
+        if x.shape[3] % 2 != 0:  # [bs, 128, Freq, Time]
+            x = torch.cat([x, x[:, :, :, -1:]], dim=3)
+        return F.interpolate(x, scale_factor=.5, mode='nearest-exact')  # F.avg_pool2d(x, 2)
+
+    def _residual(self, x):
+        x = self.actv(x)
+        x = self.conv1(x)
+        x = self.downsample_res(x)
+        x = self.actv(x)
+        x = self.conv2(x)
+        return x
+
+    def forward(self, x):
+        x = self._shortcut(x) + self._residual(x)
+        return x / math.sqrt(2)  # unit variance
+
+
+class StyleEncoder(nn.Module):
+
+    #  for both acoustic & prosodic ref_s/p
+
+    def __init__(self,
+                 dim_in=64,
+                 style_dim=128,
+                 max_conv_dim=512):
+        super().__init__()
+        blocks = [spectral_norm(nn.Conv2d(1, dim_in, 3, stride=1, padding=1))]
+        for _ in range(4):
+            dim_out = min(dim_in * 2, 
+                          max_conv_dim)
+            blocks += [ResBlk(dim_in, dim_out)]
+            dim_in = dim_out
+        blocks += [nn.LeakyReLU(0.24),  # w/o this activation - produces no speech
+                   spectral_norm(nn.Conv2d(dim_out, dim_out, 5, stride=1, padding=0)),
+                   nn.LeakyReLU(0.2)  # 0.3 sounds nice
+                   ]
+        self.shared = nn.Sequential(*blocks)
+        self.unshared = nn.Linear(dim_out, style_dim)
+
+    def forward(self, x):
+        x = self.shared(x)
+        x = x.mean(3, keepdims=True)  # comment this line for time varying style vector
+        x = x.transpose(1, 3)
+        s = self.unshared(x)
+        return s
+
+
+class LinearNorm(torch.nn.Module):
+    def __init__(self, in_dim, out_dim, bias=True):
+        super().__init__()
+        self.linear_layer = torch.nn.Linear(in_dim, out_dim, bias=bias)
+
+    def forward(self, x):
+        return self.linear_layer(x)
+
+
+class LayerNorm(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 TextEncoder(nn.Module):
+    def __init__(self, channels, kernel_size, depth, n_symbols):
+        super().__init__()
+        self.embedding = nn.Embedding(n_symbols, channels)
+        padding = (kernel_size - 1) // 2
+        self.cnn = nn.ModuleList()
+        for _ in range(depth):
+            self.cnn.append(nn.Sequential(
+                weight_norm(nn.Conv1d(channels, channels, kernel_size=kernel_size, padding=padding)),
+                LayerNorm(channels),
+                nn.LeakyReLU(0.24))
+                            )
+        self.lstm = nn.LSTM(channels, channels//2, 1,
+                            batch_first=True, bidirectional=True)
+
+    def forward(self, x):
+        x = self.embedding(x)  # [B, T, emb]
+        x = x.transpose(1, 2)
+        for c in self.cnn:
+            x = c(x)
+        x = x.transpose(1, 2)
+        x, _ = self.lstm(x)
+        return x
+
+
+class AdaLayerNorm(nn.Module):
+
+    def __init__(self, style_dim, channels=None, eps=1e-5):
+        super().__init__()
+        self.eps = eps
+        self.fc = nn.Linear(style_dim, 1024)
+
+    def forward(self, x, s):
+        h = self.fc(s)
+        gamma = h[:, :, :512]
+        beta = h[:, :, 512:1024]
+        x = F.layer_norm(x, (512, ), eps=self.eps)
+        x = (1 + gamma) * x + beta
+        return x  # [1, 75, 512]
+
+
+class ProsodyPredictor(nn.Module):
+
+    def __init__(self, style_dim, d_hid, nlayers, max_dur=50):
+        super().__init__()
+
+        self.text_encoder = DurationEncoder(sty_dim=style_dim,
+                                            d_model=d_hid,
+                                            nlayers=nlayers)  # called outside forward
+        self.lstm = nn.LSTM(d_hid + style_dim, d_hid // 2,
+                            1, batch_first=True, bidirectional=True)
+        self.duration_proj = LinearNorm(d_hid, max_dur)
+        self.shared = nn.LSTM(d_hid + style_dim, d_hid //
+                              2, 1, batch_first=True, bidirectional=True)
+        self.F0 = nn.ModuleList([
+            AdainResBlk1d(d_hid, d_hid, style_dim),
+            AdainResBlk1d(d_hid, d_hid // 2,  style_dim, upsample=True),
+            AdainResBlk1d(d_hid // 2, d_hid // 2, style_dim),
+            ])
+        self.N = nn.ModuleList([
+            AdainResBlk1d(d_hid, d_hid, style_dim),
+            AdainResBlk1d(d_hid, d_hid // 2, style_dim, upsample=True),
+            AdainResBlk1d(d_hid // 2, d_hid // 2, style_dim)
+            ])
+        self.F0_proj = nn.Conv1d(d_hid // 2, 1, 1, 1, 0)
+        self.N_proj = nn.Conv1d(d_hid // 2, 1, 1, 1, 0)
+
+    def F0Ntrain(self, x, s):
+
+        x, _ = self.shared(x)  # [bs, time, ch] LSTM
+
+        x = x.transpose(1, 2)  # [bs, ch, time]
+
+        F0 = x
+
+        for block in self.F0:
+            # print(f'LOOP {F0.shape=} {s.shape=}\n')
+            # )N F0.shape=torch.Size([1, 512, 147]) s.shape=torch.Size([1, 128])
+            # This is an AdainResBlk1d expects conv1d dimensions
+            F0 = block(F0, s)
+        F0 = self.F0_proj(F0)
+
+        N = x
+
+        for block in self.N:
+            N = block(N, s)
+        N = self.N_proj(N)
+
+        return F0, N
+    
+    def forward(self, d_en=None, s=None):
+        blend = self.text_encoder(d_en, s)
+        x, _ = self.lstm(blend)
+        dur = self.duration_proj(x)  # [bs, 150, 50]
+        
+        _, input_length, classifier_50 = dur.shape
+
+        dur = dur[0, :, :]
+        dur = torch.sigmoid(dur).sum(1)
+        dur = dur.round().clamp(min=1).to(torch.int64)
+        aln_trg = torch.zeros(1,
+                              dur.sum(),
+                              input_length, 
+                              device=s.device)
+        c_frame = 0
+        for i in range(input_length):
+            aln_trg[:, c_frame:c_frame + dur[i], i] = 1
+            c_frame += dur[i]
+        en = torch.bmm(aln_trg, blend)
+        F0_pred, N_pred = self.F0Ntrain(en, s)
+        return aln_trg, F0_pred, N_pred
+
+
+class DurationEncoder(nn.Module):
+
+    def __init__(self, sty_dim=128, d_model=512, nlayers=3):
+        super().__init__()
+        self.lstms = nn.ModuleList()
+        for _ in range(nlayers):
+            self.lstms.append(nn.LSTM(d_model + sty_dim,
+                                      d_model // 2,
+                                      num_layers=1,
+                                      batch_first=True,
+                                      bidirectional=True
+                                      ))
+            self.lstms.append(AdaLayerNorm(sty_dim, d_model))
+
+
+    def forward(self, x, style):
+
+        _, _, input_lengths = x.shape  # [bs, 512, time]
+
+        style = _tile(style, length=x.shape[2]).transpose(1, 2)
+        x = x.transpose(1, 2)
+
+        for block in self.lstms:
+            if isinstance(block, AdaLayerNorm):
+                
+                x = block(x, style)  # LSTM has transposed x
+
+            else:
+                x = torch.cat([x, style], axis=2)
+                # LSTM
+
+                x,_ = block(x)  # expects [bs, time, chan]  OUTPUTS [bs, time, 2*chan]  2x FROM BIDIRECTIONAL
+
+        return torch.cat([x, style], axis=2)  # predictor.lstm()