https://huggingface.co/Dionyssos/_TTS075B

tts.py

@@ -1,847 +1,740 @@
-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#2.9.0 cu126
+from torch import nn
 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 transformers import Wav2Vec2PreTrainedModel, PretrainedConfig#4.49.0
 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
+import re
+from collections import deque
 
-_pad = "$"
-_punctuation = ';:,.!?¡¿—…"«»“” '
-_letters = 'ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz'
-_letters_ipa = "ɑɐɒæɓʙβɔɕçɗɖðʤəɘɚɛɜɝɞɟʄɡɠɢʛɦɧħɥʜɨɪʝɭɬɫɮʟɱɯɰŋɳɲɴøɵɸθœɶʘɹɺɾɻʀʁɽʂʃʈʧʉʊʋⱱʌɣɤʍχʎʏʑʐʒʔʡʕʢǀǁǂǃˈˌːˑʼʴʰʱʲʷˠˤ˞↓↑→↗↘'̩'ᵻ"
-MAX_PHONEMES = 424   # For OOM is the max length of single (non-split) sentence for StyleTTS2 inference
+from safetensors.torch import load_file
+from sentencepiece import SentencePieceProcessor
+from einops import rearrange
 
-symbols = [_pad] + list(_punctuation) + list(_letters) + list(_letters_ipa)
 
-dicts = {}
-for i in range(len((symbols))):
-    dicts[symbols[i]] = i
 
+class ActivationGating(nn.Module):
 
-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):
+    def __init__(self, dim_feedforward=4224):
         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
-
-
+        d = 2816 if dim_feedforward == 4224 else 2048
+        self.linear_in = nn.Linear(1024, 2 * d, bias=False)
+        self.linear_out = nn.Linear(d, 1024, bias=False)
 
+    def forward(self, x):
+        x = F.linear(x, self.linear_in.weight)
+        B, T, _ = x.shape
+        x = x.view(B, T, 2, -1)
+        x = F.silu(x[:, :, 0, :]) * x[:, :, 1, :]
+        x = F.linear(x, self.linear_out.weight)
+        return x
 
-def get_padding(kernel_size, dilation=1):
-    return int((kernel_size*dilation - dilation)/2)
 
+def apply_rope(q, k, offset=0):
+    q_type = q.dtype
+    q = q.to(torch.float)
+    k = k.to(torch.float)
+    bs, h, _1, d = k.shape
 
-def _tile(x,
-          length=None):
-    x = x.repeat(1, 1, int(length / x.shape[2]) + 1)[:, :, :length]
-    return x
+    # fr = torch.exp(-18.420680743952367 / d * torch.arange(d // 2, device=q.device, dtype=torch.float))
+    # fr = torch.exp(-18.42068099975586 / d * torch.arange(d // 2, device=q.device, dtype=torch.float))
+    fr = torch.exp(-18.4206809997 / d * torch.arange(d // 2, device=q.device, dtype=torch.float))
 
+    t = offset * fr[None, None, :, None]
 
-class AdaIN1d(nn.Module):
+    r = torch.cos(t)
+    i = torch.sin(t)
 
-    # used by HiFiGan & ProsodyPredictor
+    q = q.view(bs, h, d // 2, 2)  # interleave
+    k = k.view(bs, h, d // 2, 2)
 
-    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
+    qor = q[:, :, :, :1] * r - q[:, :, :, 1:] * i
+    qoi = q[:, :, :, :1] * i + q[:, :, :, 1:] * r
+    kor = k[:, :, :, :1] * r - k[:, :, :, 1:] * i
+    koi = k[:, :, :, :1] * i + k[:, :, :, 1:] * r
 
+    qo = torch.cat([qor.to(dtype=q_type), qoi.to(dtype=q_type)], dim=3)
+    ko = torch.cat([kor.to(dtype=q_type), koi.to(dtype=q_type)], dim=3)
 
-class SourceModuleHnNSF(torch.nn.Module):
+    return qo.view(bs, h, 1, d), ko.view(bs, h, 1, d)
 
-    def __init__(self):
 
+class RMSNorm(nn.Module):
+    def __init__(self, d=1024):
         super().__init__()
-        self.harmonic_num = 8
-        self.l_linear = torch.nn.Linear(self.harmonic_num + 1, 1)
-        self.upsample_scale = 300
-        
+        self.alpha = nn.Parameter(torch.full((1, 1, d), 1.0, dtype=torch.float64))
 
     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()
+        x = x.to(torch.float64)
+        v = 9e-9 + torch.mean(x * x, dim=2, keepdim=True)
+        return (x * (self.alpha * torch.rsqrt(v))).to(torch.bfloat16)
 
-        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)))
+class LLMAttention(nn.Module):
 
-            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 __init__(self, weights_per_step):
+        super().__init__()
+        self.weights_per_step = weights_per_step
+        self.k_history = None
+        self.v_history = None
+        p = 9 if weights_per_step else 1
+        self.out_projs = nn.ModuleList([nn.Linear(1024, 1024, bias=False) for _ in range(p)])
+        self.in_projs = nn.ModuleList([nn.Linear(1024, 3 * 1024, bias=False) for _ in range(p)])
 
-    def forward(self, x, s, f0):
+    def forward(self, query):
 
-        # 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)
+        offset = 0 if self.k_history is None else self.k_history.shape[2]  # if overpass RoPE untrained or DPF 16x
 
-        # x.shape=torch.Size([1, 512, 484]) s.shape=torch.Size([1, 1, 1, 128]) f0.shape=torch.Size([1, 145200, 1]) GENERAT 253
+        if (self.weights_per_step and offset % self.weights_per_step == 0) or (offset % 473 == 0):
+            self.k_history = None
+            self.v_history = None
+            offset = 0
 
-        # [1, 145400, 1] f0 enters already upsampled to full wav 24kHz length
-        har_source = self.m_source(f0)
+        if self.weights_per_step:
+            x = self.in_projs[offset if offset < 9 else 8](query)
+        else:
+            x = self.in_projs[0](query)
+        q, k, v = rearrange(x, "b t (p h d) -> p b h t d", p=3, h=16)
+        q, k = apply_rope(q, k, offset=offset)        
+        # KVCACHE
+        if self.k_history is not None:
+            self.k_history = torch.cat([self.k_history, k], 2)
+            self.v_history = torch.cat([self.v_history, v], 2)
+        else:
+            self.k_history = k
+            self.v_history = v            
+        k = self.k_history
+        v = self.v_history
+        # ones-bool attn mask sounds better than passing no mask argument
+        x = F.scaled_dot_product_attention(q, k, v, torch.ones(k.shape[0], 1, 1, k.shape[2],dtype=torch.bool, device=k.device))
+        x = rearrange(x, "b h t d -> b t (h d)")
+        if self.weights_per_step:
+            return self.out_projs[offset if offset < 9 else 8](x)
+        return self.out_projs[0](x)
 
-        har_source = har_source.transpose(1, 2)
 
-        for i in range(self.num_upsamples):
+class LLMTransformerLayer(nn.Module):
 
-            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)
+    def __init__(self, weights_per_step=None):
+        super().__init__()
+        self.self_attn = LLMAttention(weights_per_step=weights_per_step)
+        self.norm1 = RMSNorm()
+        self.norm2 = RMSNorm()
+        self.weights_per_step = weights_per_step
+        if self.weights_per_step:
+            self.gating = nn.ModuleList([ActivationGating(3072) for _ in range(9)])
+        else:
+            self.gating = ActivationGating()
 
-            x = self.ups[i](x)
+    def forward(self, x):
+        x = self.self_attn(self.norm1(x)) + x
+        if self.weights_per_step:
+            p = self.self_attn.k_history.shape[2] - 1
+            return x + self.gating[p if p < 9 else 8](self.norm2(x))
+        return x + self.gating(self.norm2(x))
 
-            x = x + x_source
 
-            xs = None
-            for j in range(self.num_kernels):
+class LLMTransformer(nn.Module):
 
-                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)
+    def __init__(
+        self,
+        num_layers=24,
+        weights_per_step=False):
+        super().__init__()
+        self.layers = nn.ModuleList(
+            [
+                LLMTransformerLayer(weights_per_step=weights_per_step)
+                for _ in range(num_layers)
+                ])
 
+    def forward(self, x):
+        for lay in self.layers:
+                x = lay(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()
+class Voc(Wav2Vec2PreTrainedModel):
+
+    '''For using different batch_siz -> Voc._flush()
+    '''
+
+    def __init__(self, config=PretrainedConfig()):
+        super().__init__(config=config)
+        self.encoder_transformer = VocTransformer()
+        self.decoder_transformer = VocTransformer()
+        self.encoder = SEANetEncoder()
+        self.decoder = SEANetDecoder()
+        self.sample_rate = 24000
+        self.quantizer = SplitResidualVectorQuantizer()
+        self.downsample = BufferConv1d(512, 512, kernel_size=4, stride=2, groups=1, bias=False)
+        upsample_channel_wise_bug = True
+        self.upsample = BufferConvTranspose1d(512, 512, kernel_size=4,
+                                              groups=512 if upsample_channel_wise_bug else 1,
+                                              stride=2, bias=False)
+        self.frame_rate = 12.5
+        self.encode_buffer = None
+
+    def _flush(self):
+        '''stream buffers have tensors of old batch size! Voc()._flush() to clean buffers
+        '''
+        self.encode_buffer = None # holds unused (incomplete windows of len < 1920) - we need 1920 to produce 1 token
+        if self.downsample.previous is not None:
+            self.downsample.previous = None
+        if self.upsample.partial is not None:
+            self.upsample.partial = None
+        for arch in [self.encoder, self.decoder]:
+            for _m in arch.model:
+                if type(_m) is SEANetResnetBlock:
+                    for _b in _m.block:
+                        if type(_b) is BufferConv1d:
+                            if _b.previous is not None:
+                                _b.previous = None
+                if type(_m) is BufferConv1d:
+                    if _m.previous is not None:
+                        _m.previous = None
+                if type(_m) is BufferConvTranspose1d:
+                    if _m.partial is not None:
+                        _m.partial = None
+
+    @torch.no_grad()
+    def encode(self, x):
+        '''24KHz audio to codes
+           x : [bs, 1, 24 KHz]
+           c : [bs, 8, time]      = 1920 audio samples produce 1 time frame (of n_q codebooks)
+        '''
+        if self.encode_buffer is not None:
+            x = torch.cat([self.encode_buffer, x], 2)
+        _bs, _1, _len = x.shape
+        num_frames = int(_len / 1920)
+        leftover = x[:, :, (num_frames+1) * 1920:]
+        if leftover.shape[2] > 0:
+            self.encode_buffer = leftover
         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
+            self.encode_buffer = None
+            torch.cuda.empty_cache()
+        if num_frames > 0:
+            c = []
+            for n in range(num_frames):
+                e = self.encoder(x[:, :, n * 1920:(n + 1) * 1920])
+                e = self.encoder_transformer(e)
+                e = self.downsample(e)
+                _c = self.quantizer.encode(e)
+                c.append(_c)
+            c = torch.cat(c, 2)
+        else:
+            # num_frames = 0  Early exit -> for x.shape[2]<1920 fill conv buffers but can't output token
+            c = torch.empty(_bs, 16, 0)
+        return c
+
+    @torch.no_grad()
+    def decode(self, c):
+        '''codes to 24kHZ audio
+           c: [bs, 8, n_tokens]
+           x: [bs, 1, n_tokens * 1920]
+        '''
+        _hidden = []
+        for i in range(c.shape[2]):
+            x = self.quantizer.decode(c[:, :, i:i+1])
+            x = self.upsample(x)
+            x = self.decoder_transformer(x)
+            x = self.decoder(x)
+            _hidden.append(x)
+        return torch.cat(_hidden, 2)  # [bs, 1, 24KHz]
+
+
+class SEANetResnetBlock(nn.Module):
+    def __init__(
+        self,
+        dim,
+        kernel_sizes=[3, 1],
+    ):
+        super().__init__()
 
-    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
+        block = []
+        for i, kernel_size in enumerate(kernel_sizes):
 
-    def forward(self, x, s):
-        out = self._residual(x, s)
-        out = (out + self._shortcut(x)) / math.sqrt(2)
-        return out
+            block += [
+                nn.ELU(),
+                BufferConv1d(
+                    dim      if i == 0 else dim // 2,
+                    dim // 2 if i == 0 else dim,
+                    kernel_size=kernel_size,
+                    bias=True,
+                ),
+            ]
 
+        self.block = nn.Sequential(*block)
 
-class UpSample1d(nn.Module):
-    def __init__(self, layer_type):
+    def forward(self, x):
+        return x + self.block(x)
+
+
+class SEANetEncoder(nn.Module):
+    def __init__(
+        self,
+        channels=1, # DOES NOT SUPPORT STEREO
+        dimension=512,
+        n_filters=64,
+        ratios=[8, 6, 5, 4],
+        kernel_size=7,
+        last_kernel_size=3,
+    ):
         super().__init__()
-        self.layer_type = layer_type
+        self.ratios = list(reversed(ratios))
+        del ratios
+        mult = 1
+        model=[
+            BufferConv1d(
+                channels,
+                mult * n_filters,
+                kernel_size,
+                bias=True
+            )
+        ]
+        for i, ratio in enumerate(self.ratios):
+            model += [SEANetResnetBlock(mult * n_filters),
+                      nn.ELU(),
+                      BufferConv1d(mult * n_filters,
+                                      mult * n_filters * 2,
+                                      kernel_size=ratio * 2,
+                                      stride=ratio,
+                                      bias=True)]
+            mult *= 2
+        # ENDFOR
+        model += [nn.ELU(),
+                  BufferConv1d(mult * n_filters,
+                                    dimension,
+                                    last_kernel_size,
+                                    bias=True)]
+        self.model = nn.Sequential(*model)
 
     def forward(self, x):
-        if self.layer_type == 'none':
-            return x
-        else:
-            return F.interpolate(x, scale_factor=2, mode='nearest-exact')
-
+        return self.model(x)
 
-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()
+class SEANetDecoder(nn.Module):
 
-        self.encode = AdainResBlk1d(dim_in + 2, 1024, style_dim)
+    def __init__(
+        self,
+        channels=1,
+        dimension=512,
+        n_filters=64,
+        ratios=[8, 6, 5, 4],
+        kernel_size=7,
+        last_kernel_size=3):
 
-        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))
+        super().__init__()
+        mult = int(2 ** len(ratios))
+        model = [BufferConv1d(dimension,
+                                 mult * n_filters,
+                                 kernel_size,
+                                 bias=True)]
+        #UP
+        for i, ratio in enumerate(ratios):
+            model += [nn.ELU(),
+                      BufferConvTranspose1d(mult * n_filters,
+                                        mult * n_filters // 2,
+                                        kernel_size=ratio * 2,
+                                        stride=ratio,
+                                        bias=True),
+                      SEANetResnetBlock(mult * n_filters // 2)]
+            mult //= 2
+        # LAST
+        model += [
+            nn.ELU(),
+            BufferConv1d(
+                n_filters,
+                channels,
+                last_kernel_size,
+                bias=True
+            ),
+        ]
+        self.model = nn.Sequential(*model)
 
-        self.F0_conv = weight_norm(
-            nn.Conv1d(1, 1, kernel_size=3, stride=2, groups=1, padding=1))  # smooth
+    def forward(self, x):
+        return self.model(x)
 
-        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)),
-        )
+class BufferConv1d(nn.Conv1d):
+    def __init__(self,
+                 *args,
+                 **kwargs):
+        super().__init__(*args, **kwargs)
+        self.previous = None
 
-        self.generator = Generator(style_dim, resblock_kernel_sizes, upsample_rates,
-                                   upsample_initial_channel, resblock_dilation_sizes, upsample_kernel_sizes)
+    def forward(self, x):
+        k = self.kernel_size[0]
 
-    def forward(self, asr=None, F0_curve=None, N=None, s=None):
+        if self.previous is not None:
 
+                x = torch.cat([self.previous, x], 2)
 
-        F0 = self.F0_conv(F0_curve)
-        N = self.N_conv(N)
+        else:  # If self.previous is None => Use zero pad
 
+            if k == 3:
 
-        x = torch.cat([asr, F0, N], axis=1)
+                p = (2, 0)
+                x = F.pad(x, p, mode='replicate', value=0.0) # skip connections SeaNetResBlk
 
-        x = self.encode(x, s)
+            elif k == 4:  # ConvTrUpsample is the first conv encountered by decode replicate solves pulse
 
-        asr_res = self.asr_res(asr)
+                p = (3, 0)
+                x = F.pad(x, p, mode='replicate', value=0.0)
 
-        res = True
-        for block in self.decode:
-            if res:
+            elif k == 7:
 
-                x = torch.cat([x, asr_res, F0, N], axis=1)
+                p = (6, 0)
+                x = F.pad(x, p, mode='replicate', value=0.0)
 
-            x = block(x, s)
-            if block.upsample_type != "none":
-                res = False
+            elif k == 16:
 
-        x = self.generator(x, s, F0_curve)
-        return x
+                p = (2, 0)
+                x = F.pad(x, p, mode='replicate', value=0.0)  # THis can be also constant w/o pulse occur
 
+        num_frames = int(         (x.shape[2] - self.kernel_size[0])  / self.stride[0] )  +  1  # +1 is: k starts at left of x and doing (I-k)/s jumps
+        offset = num_frames * self.stride[0]
+        self.previous = x[..., offset:]
+        return super().forward(x)
 
-class MelSpec(torch.nn.Module):
 
+class BufferConvTranspose1d(nn.ConvTranspose1d):
+    # kernel 5 has only 1 pixel for input (cloned)
+    # https://distill.pub/2016/deconv-checkerboard/
     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))
+                 *args,
+                 **kwargs):
+        super().__init__(*args,
+                         **kwargs)
+        self.partial = None
         
     def forward(self, x):
-        return self.conv(x)
+        out = super().forward(x)
+        OT = out.shape[2]
+        invalid_steps = self.kernel_size[0] - self.stride[0]
+        if self.partial is not None:
+            PT = self.partial.shape[-1]
+            if self.bias is not None:
+                out[..., :PT] += self.partial - self.bias[:, None]
+            else:
+                out[..., :PT] += self.partial  # for ConvTrUpsample1d
+        invalid_steps = self.kernel_size[0] - self.stride[0]
+        self.partial = out[..., OT - invalid_steps :]
+        out = out[...,:OT - invalid_steps]
+        return out
 
 
-class ResBlk(nn.Module):
-    def __init__(self, 
-                 dim_in, dim_out):
+class CodeBook(nn.Module):
+    def __init__(self, dim, codebook_size):
         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
+        self.register_buffer('_e', torch.zeros(codebook_size, dim))
 
-    def forward(self, x):
-        x = self._shortcut(x) + self._residual(x)
-        return x / math.sqrt(2)  # unit variance
+    def encode(self, x):
+        dist = torch.cdist(
+            x.transpose(1, 2), # [bs, time, 256]
+            self._e[None, :, :]  # [1, 2048, 256]
+        )
+        codes = dist.argmin(2)
+        return codes
 
+    def decode(self, codes):
+        quantized = F.embedding(codes, self._e)
+        return quantized.transpose(1, 2)  # [1, 256, time]
 
-class StyleEncoder(nn.Module):
 
-    #  for both acoustic & prosodic ref_s/p
+class SplitResidualVectorQuantizer(nn.Module):
 
     def __init__(self,
-                 dim_in=64,
-                 style_dim=128,
-                 max_conv_dim=512):
+                 n_q=None):
         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
+        self.in_proj_s  = torch.nn.Conv1d(512, 256, 1, bias=False)
+        self.in_proj_a  = torch.nn.Conv1d(512, 256, 1, bias=False)
+        self.out_proj_s = torch.nn.Conv1d(256, 512, 1, bias=False) # reused for all _acoustic_books
+        self.out_proj_a = torch.nn.Conv1d(256, 512, 1, bias=False)
+        self.layers = nn.ModuleList([CodeBook(dim=256, codebook_size=2048) for _ in range(18)])
+        self._acoustic_books = range(1, 16)  # Official Mimi
+        # CODEBOOKS
+        # Here we re use RVQ codebooks for higher fidelity!
+        # Exclude 0 here as it has different proj (in_proj_s)
+        # self._acoustic_books = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 17, 17, 17, 17]
+
+    def encode(self, x):
+        indices = self.layers[0].encode(self.in_proj_s(x))  # integers
+        all_indices = [ indices[:, None, :], ]
+        x = self.in_proj_a(x)
+        for _cb in self._acoustic_books:
+            indices = self.layers[_cb].encode(x)
+            x = x - self.layers[_cb].decode(indices)
+            all_indices.append(indices[:, None, :])
+        codes = torch.cat(all_indices, 1)
+        return codes
+
+    def decode(self, codes):
+        _s = self.layers[0].decode(codes[:, 0, :])
+        _a = torch.zeros([1, 1], device=codes.device)
+        for i, _cb in enumerate(self._acoustic_books):
+            _a = _a + self.layers[_cb].decode(codes[:, i+1, :])
+        return self.out_proj_s(_s) + self.out_proj_a(_a)   # [bs, 512, time]
+
+
+class VocAttention(nn.Module):
 
+    def __init__(self,
+                 embed_dim):
 
-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)
+        self.fused_proj = nn.Parameter(torch.zeros(embed_dim, embed_dim))
 
     def forward(self, x):
-        return self.linear_layer(x)
-
+        '''bypass of streaming training'''
+        if x.shape[1] > 1:
+            x = x.mean(1, keepdims=True)
+        x = torch.matmul(x, self.fused_proj)
+        return x  # FFN broadcasts to x.shape[1]=2
 
-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 VocTransformerLayer(nn.Module):
 
-class TextEncoder(nn.Module):
-    def __init__(self, channels, kernel_size, depth, n_symbols):
+    def __init__(self, d_model=512, dim_feedforward=2048):
         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)
+        self.self_attn = VocAttention(embed_dim=d_model)
+        self.norm1 = nn.LayerNorm(d_model, eps=1e-5)
+        self.norm2 = nn.LayerNorm(d_model, eps=1e-5)
+        self.linear1 = nn.Linear(d_model, dim_feedforward, bias=False)
+        self.linear2 = nn.Linear(dim_feedforward, d_model, bias=False)
 
     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
+        x = x + self.self_attn(self.norm1(x))
+        return x + self.linear2(F.gelu(self.linear1(self.norm2(x))))
 
 
-class AdaLayerNorm(nn.Module):
+class VocTransformer(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__()
+    def __init__(self):
 
-        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):
+        self.layers = nn.ModuleList(VocTransformerLayer() for _ in range(8))
 
-        _, _, input_lengths = x.shape  # [bs, 512, time]
-
-        style = _tile(style, length=x.shape[2]).transpose(1, 2)
+    def forward(self, x):
         x = x.transpose(1, 2)
+        for la in self.layers:
+            x = la(x)
+        return x.transpose(1, 2)
 
-        for block in self.lstms:
-            if isinstance(block, AdaLayerNorm):
-                
-                x = block(x, style)  # LSTM has transposed x
+class Entry():
+    def __init__(self, tokens=None):
+        self.tokens = tokens
+        self.padding = len(tokens) + 2 - 1
 
-            else:
-                x = torch.cat([x, style], axis=2)
-                # LSTM
+class TokenState:
 
-                x,_ = block(x)  # expects [bs, time, chan]  OUTPUTS [bs, time, 2*chan]  2x FROM BIDIRECTIONAL
+    def __init__(self, entries = None):
+        self.entries = entries
+        self.queued = deque([])
+        self.lookahead_queued = deque()
+        self.end_step = None
+        self.forced_padding = 2
 
-        return torch.cat([x, style], axis=2)  # predictor.lstm()
+class TTSModel(nn.Module):
+
+    def __init__(self):
+        super().__init__()
+        self.tokenizer = SentencePieceProcessor(str(hf_hub_download(repo_id='kyutai/tts-0.75b-en-public',
+                                                                    filename='tokenizer_spm_8k_en_fr_audio.model')))
+        with torch.device("meta"):
+            self.emb = nn.ModuleList([ScaledEmbedding(2049, 1024) for _ in range(16)])
+            self.text_emb = ScaledEmbedding(8001, 1024, demux_second_stream=True)
+            self.transformer = LLMTransformer()
+            self.out_norm = RMSNorm()
+            self.depformer_in = nn.ModuleList([nn.Linear(1024, 1024, bias=False) for _ in range(9)])
+            self.depformer_emb = nn.ModuleList([ScaledEmbedding(2049, 128) for _ in range(16 - 1)])
+            self.depformer_text_emb = ScaledEmbedding(8001, 128, demux_second_stream=True)
+            self.depformer = LLMTransformer(num_layers=4, weights_per_step=16)
+            self.linears = nn.ModuleList([nn.Linear(1024, 2048, bias=False) for _ in range(16)])  # DPF heads
+
+        state_d = load_file(hf_hub_download(repo_id='Dionyssos/_TTS075B', filename='tts_075B.safetensors'))
+        self.load_state_dict(state_d, assign=True, strict=True)  #overwrite devices of rand init params
+        self.to(dtype=torch.bfloat16).eval()
+
+    def prepare_script(self, script='Type your text here.'):
+        entries = []
+        # break is indicated as e.g. <break time="3s"/>
+        event_re = re.compile(r"(?:<break\s+time=\"([0-9]+(?:.[0-9]*)?)s\"\s*/?>)|(?:\s+)")
+        line = script.replace('’', "'").replace(':', " ").replace('(', "").replace(')', "")
+        while line:
+            match = event_re.search(line)
+            if match is None:
+                break
+            word = line[:match.start()]
+            line = line[match.end():]
+            if word:
+                entries.append(Entry(tokens=self.tokenizer.encode(word)))
+            if match.group(1):
+                raise ValueError
+                # break_duration = float(match.group(1))
+                # padding = int(round(break_duration * frame_rate))
+                # entry = Entry(tokens=[], text='', padding=padding)
+                # entries.append(entry)
+        if line:
+            entries.append(Entry(tokens=self.tokenizer.encode(line)))
+        return entries
+
+    @property
+    def device(self):
+        return next(iter(self.parameters())).device
+
+    @torch.no_grad()
+    def generate(self, text=None,
+                 _wav=None, mimi=None,
+                 play=16):
+
+        state = TokenState(entries=deque(self.prepare_script(script=text)))
+        upper_lim = 9999
+        self.cache = torch.full((2,17, 4), -1, device=self.device, dtype=torch.long)
+        pcms = []#final audio to return
+        for offset in range(upper_lim):
+            print(f'{offset=} of {upper_lim=}',end='\r')
+            if state.end_step is not None:
+                if offset >= state.end_step + 16 + 4:
+                    break
+
+            input_ = self.cache[:, :, offset % self.cache.shape[2]].clone()
+
+            if offset == 0:
+                input_[:, 0] = 8000  # so we dont have to reset cfg txr = -1 for offset >0
+                input_[:, 1:] = 2048
+
+            if offset < 3:
+                input_[:, 2:] = 2048
+
+
+            x = self.text_emb(input_[:, :1])
+            for cb_ in range(16):
+                x = self.emb[cb_](input_[:, cb_ + 1 : cb_ + 2]) + x
+            x = self.out_norm(self.transformer(x))
+
+
+            token = -1
+            if offset > _wav.shape[1]:
+                token = 0
+                # START
+                if state.queued:
+                    token = 3
+                if state.forced_padding > 0:
+                    token = 3
+                #===================================
+                if token == 0:
+                    if state.entries:
+                        e = state.entries.popleft()
+                        if e.tokens:
+                            state.queued.extend(e.tokens)
+                            lookahead =2
+                            for e2 in state.entries:
+                                if e2.tokens:
+                                    lookahead -= 1
+                                    if lookahead == 0:
+                                        state.lookahead_queued.extend(e2.tokens)
+                                        break
+                            # print('\neeee',e2,'\n\n')
+                            # raise ValueError
+                        else:
+                            token = 3
+                        state.forced_padding = e.padding
+                        # print(f'\n\n=========o=============\n{state.lookahead_queued=} {state.queued=}===================\n\n')
+                    else:
+                        token = 3
+                        if state.end_step is None:
+                            token = 0
+                        if state.end_step is None:
+                            state.end_step = offset
+                #==============================================
+                output=0
+                if token == 3:
+                    if state.forced_padding > 0:
+                        state.forced_padding -= 1
+                    if state.queued:
+                        output = state.queued.popleft()
+                    else:
+                        output = 3
+                # ==========================
+                second = -1
+                if output == 0:
+                    second = 0
+                    if state.queued:
+                        output = state.queued.popleft()
+                    else:
+                        output = 3
+                elif state.lookahead_queued:
+                    second = state.lookahead_queued.popleft()  # Difference of queued and lookahead_queued?
+                token = (second + 1) * 8001 + output
+
+            # audio tokens
+            ac = (offset + 1) % self.cache.shape[2]
+            self.cache[0, 0, ac] = token
+            audio_tokens = torch.ones([1, 16], device=x.device, dtype=torch.long)
+            if offset > play:
+                prev_token = torch.tensor([[token]], device=x.device, dtype=torch.long)
+                for _cb in range(16):
+                    last_token_input = None
+                    if _cb == 0:
+                        last_token_input = self.depformer_text_emb(prev_token.repeat(2, 1))
+                    else:
+                        last_token_input = self.depformer_emb[_cb - 1](prev_token)
+                    dep_output = self.depformer(self.depformer_in[_cb if _cb < 9 else 8](x) + last_token_input)
+                    logits = self.linears[_cb](dep_output)
+                    prev_token = (2.0 * logits[0, :, :] - logits[1, :, :]).argmax(1)
+                    audio_tokens[0, _cb] = prev_token
+                if offset > play and offset < play + 1 + _wav.shape[1]:
+                     audio_tokens[:, :5] = _wav[:5, offset - play - 3]
+                     audio_tokens[:, 11:] = _wav[11:, offset - play - 3]
+                # next turn
+                self.cache[0, 1:, ac] = audio_tokens
+                # cfg
+                if offset > 16 + 2 + _wav.shape[1]:
+                    if offset > 16 + 4 + _wav.shape[1]:
+                        self.cache[1, 1:, ac] = self.cache[0, 1:, ac]
+                    else:
+                        self.cache[1, 1, ac] = self.cache[0, 1, ac]
+            # ivao0/voc
+            if offset > 20 + _wav.shape[1]:
+                audio_tokens[:, 0] = self.cache[0, 1, (offset - 1) % self.cache.shape[2]]  # previous
+                pcms.append(mimi.decode(audio_tokens[:, :, None]))  # [1,1,1920]
+        x = torch.cat(pcms, dim=2)[0, 0, :]
+        return x.cpu().numpy()
+
+    def _flush(self):
+        for lay in self.transformer.layers:
+            lay.self_attn.k_history = None
+            lay.self_attn.v_history = None
+
+
+class ScaledEmbedding(nn.Embedding):
+    def __init__(self, num_embeddings=None, embedding_dim=None, demux_second_stream=False):
+        super().__init__(num_embeddings, embedding_dim)
+        self.zero_idx = -1
+        self.low_rank = None
+        self.demux_second_stream = demux_second_stream
+        if self.demux_second_stream:
+            self.out1 = nn.Linear(embedding_dim, 1024, bias=False)
+            self.out2 = nn.Linear(embedding_dim, 1024, bias=False)
+        else:
+            if embedding_dim != 1024:
+                self.low_rank = nn.Linear(embedding_dim, 1024, bias=False)
+
+    def forward(self, input):
+        is_zero = input == self.zero_idx
+        zero = torch.zeros(1, dtype=input.dtype, device=input.device)
+        input = input.clamp(min=0)
+        if self.demux_second_stream:
+            left = super().forward(input % self.num_embeddings)
+            right = input // self.num_embeddings - 1
+            right_zero = (right < 0)[..., None]
+            right.clamp_(min=0)
+            right = super().forward(right)
+            y = self.out1(left) + torch.where(right_zero, zero, self.out2(right))
+            y = torch.where(is_zero[..., None], zero, y)
+        else:
+            y = super().forward(input)
+            y = torch.where(is_zero[..., None], zero, y)
+            if self.low_rank is not None:
+                # Can only see low_rank if no demux second stream
+                y = self.low_rank(y)   # applies after
+        return y