import torch#2.9.0 cu126
from torch import nn
import torch.nn.functional as F
from transformers import Wav2Vec2PreTrainedModel, PretrainedConfig#4.49.0
from huggingface_hub import hf_hub_download
import re
from collections import deque
from safetensors.torch import load_file
from sentencepiece import SentencePieceProcessor
from einops import rearrange
class ActivationGating(nn.Module):
def __init__(self, dim_feedforward=4224):
super().__init__()
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 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
# 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]
r = torch.cos(t)
i = torch.sin(t)
q = q.view(bs, h, d // 2, 2) # interleave
k = k.view(bs, h, d // 2, 2)
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)
return qo.view(bs, h, 1, d), ko.view(bs, h, 1, d)
class RMSNorm(nn.Module):
def __init__(self, d=1024):
super().__init__()
self.alpha = nn.Parameter(torch.full((1, 1, d), 1.0, dtype=torch.float64))
def forward(self, x):
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)
class LLMAttention(nn.Module):
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, query):
offset = 0 if self.k_history is None else self.k_history.shape[2] # if overpass RoPE untrained or DPF 16x
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
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)
class LLMTransformerLayer(nn.Module):
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()
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))
class LLMTransformer(nn.Module):
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 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.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__()
block = []
for i, kernel_size in enumerate(kernel_sizes):
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)
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.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):
return self.model(x)
class SEANetDecoder(nn.Module):
def __init__(
self,
channels=1,
dimension=512,
n_filters=64,
ratios=[8, 6, 5, 4],
kernel_size=7,
last_kernel_size=3):
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)
def forward(self, x):
return self.model(x)
class BufferConv1d(nn.Conv1d):
def __init__(self,
*args,
**kwargs):
super().__init__(*args, **kwargs)
self.previous = None
def forward(self, x):
k = self.kernel_size[0]
if self.previous is not None:
x = torch.cat([self.previous, x], 2)
else: # If self.previous is None => Use zero pad
if k == 3:
p = (2, 0)
x = F.pad(x, p, mode='replicate', value=0.0) # skip connections SeaNetResBlk
elif k == 4: # ConvTrUpsample is the first conv encountered by decode replicate solves pulse
p = (3, 0)
x = F.pad(x, p, mode='replicate', value=0.0)
elif k == 7:
p = (6, 0)
x = F.pad(x, p, mode='replicate', value=0.0)
elif k == 16:
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 BufferConvTranspose1d(nn.ConvTranspose1d):
# kernel 5 has only 1 pixel for input (cloned)
# https://distill.pub/2016/deconv-checkerboard/
def __init__(self,
*args,
**kwargs):
super().__init__(*args,
**kwargs)
self.partial = None
def forward(self, 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 CodeBook(nn.Module):
def __init__(self, dim, codebook_size):
super().__init__()
self.register_buffer('_e', torch.zeros(codebook_size, dim))
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 SplitResidualVectorQuantizer(nn.Module):
def __init__(self,
n_q=None):
super().__init__()
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):
super().__init__()
self.fused_proj = nn.Parameter(torch.zeros(embed_dim, embed_dim))
def forward(self, 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 VocTransformerLayer(nn.Module):
def __init__(self, d_model=512, dim_feedforward=2048):
super().__init__()
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 = x + self.self_attn(self.norm1(x))
return x + self.linear2(F.gelu(self.linear1(self.norm2(x))))
class VocTransformer(nn.Module):
def __init__(self):
super().__init__()
self.layers = nn.ModuleList(VocTransformerLayer() for _ in range(8))
def forward(self, x):
x = x.transpose(1, 2)
for la in self.layers:
x = la(x)
return x.transpose(1, 2)
class Entry():
def __init__(self, tokens=None):
self.tokens = tokens
self.padding = len(tokens) + 2 - 1
class TokenState:
def __init__(self, entries = None):
self.entries = entries
self.queued = deque([])
self.lookahead_queued = deque()
self.end_step = None
self.forced_padding = 2
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.
event_re = re.compile(r"(?:)|(?:\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