Views
No views yet
1import torch#2.9.0 cu126
2from torch import nn
3import torch.nn.functional as F
4from transformers import Wav2Vec2PreTrainedModel, PretrainedConfig#4.49.0
5from huggingface_hub import hf_hub_download
6import re
7from collections import deque
8import sphn
9from safetensors.torch import load_file
10from sentencepiece import SentencePieceProcessor
11from einops import rearrange
12
13
14
15class ActivationGating(nn.Module):
16
17 def __init__(self, dim_feedforward=4224):
18 super().__init__()
19 d = 2816 if dim_feedforward == 4224 else 2048
20 self.linear_in = nn.Linear(1024, 2 * d, bias=False)
21 self.linear_out = nn.Linear(d, 1024, bias=False)
22
23 def forward(self, x):
24 x = F.linear(x, self.linear_in.weight)
25 B, T, _ = x.shape
26 x = x.view(B, T, 2, -1)
27 x = F.silu(x[:, :, 0, :]) * x[:, :, 1, :]
28 x = F.linear(x, self.linear_out.weight)
29 return x
30
31
32def apply_rope(q, k, offset=0):
33 q_type = q.dtype
34 q = q.to(torch.float)
35 k = k.to(torch.float)
36 bs, h, _1, d = k.shape
37
38 # fr = torch.exp(-18.420680743952367 / d * torch.arange(d // 2, device=q.device, dtype=torch.float))
39 # fr = torch.exp(-18.42068099975586 / d * torch.arange(d // 2, device=q.device, dtype=torch.float))
40 fr = torch.exp(-18.4206809997 / d * torch.arange(d // 2, device=q.device, dtype=torch.float))
41
42 t = offset * fr[None, None, :, None]
43
44 r = torch.cos(t)
45 i = torch.sin(t)
46
47 q = q.view(bs, h, d // 2, 2) # interleave
48 k = k.view(bs, h, d // 2, 2)
49
50 qor = q[:, :, :, :1] * r - q[:, :, :, 1:] * i
51 qoi = q[:, :, :, :1] * i + q[:, :, :, 1:] * r
52 kor = k[:, :, :, :1] * r - k[:, :, :, 1:] * i
53 koi = k[:, :, :, :1] * i + k[:, :, :, 1:] * r
54
55 qo = torch.cat([qor.to(dtype=q_type), qoi.to(dtype=q_type)], dim=3)
56 ko = torch.cat([kor.to(dtype=q_type), koi.to(dtype=q_type)], dim=3)
57
58 return qo.view(bs, h, 1, d), ko.view(bs, h, 1, d)
59
60
61class RMSNorm(nn.Module):
62 def __init__(self, d=1024):
63 super().__init__()
64 self.alpha = nn.Parameter(torch.full((1, 1, d), 1.0, dtype=torch.float64))
65
66 def forward(self, x):
67 x = x.to(torch.float64)
68 v = 9e-9 + torch.mean(x * x, dim=2, keepdim=True)
69 return (x * (self.alpha * torch.rsqrt(v))).to(torch.bfloat16)
70
71
72class LLMAttention(nn.Module):
73
74 def __init__(self, weights_per_step):
75 super().__init__()
76 self.weights_per_step = weights_per_step
77 self.k_history = None
78 self.v_history = None
79 p = 9 if weights_per_step else 1
80 self.out_projs = nn.ModuleList([nn.Linear(1024, 1024, bias=False) for _ in range(p)])
81 self.in_projs = nn.ModuleList([nn.Linear(1024, 3 * 1024, bias=False) for _ in range(p)])
82
83 def forward(self, query):
84
85 offset = 0 if self.k_history is None else self.k_history.shape[2] # if overpass RoPE untrained or DPF 16x
86
87 if (self.weights_per_step and offset % self.weights_per_step == 0) or (offset % 473 == 0):
88 self.k_history = None
89 self.v_history = None
90 offset = 0
91
92 if self.weights_per_step:
93 x = self.in_projs[offset if offset < 9 else 8](query)
94 else:
95 x = self.in_projs[0](query)
96 q, k, v = rearrange(x, "b t (p h d) -> p b h t d", p=3, h=16)
97 q, k = apply_rope(q, k, offset=offset)
98 # KVCACHE
99 if self.k_history is not None:
100 self.k_history = torch.cat([self.k_history, k], 2)
101 self.v_history = torch.cat([self.v_history, v], 2)
102 else:
103 self.k_history = k
104 self.v_history = v
105 k = self.k_history
106 v = self.v_history
107 # ones-bool attn mask sounds better than passing no mask argument
108 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))
109 x = rearrange(x, "b h t d -> b t (h d)")
110 if self.weights_per_step:
111 return self.out_projs[offset if offset < 9 else 8](x)
112 return self.out_projs[0](x)
113
114
115class LLMTransformerLayer(nn.Module):
116
117 def __init__(self, weights_per_step=None):
118 super().__init__()
119 self.self_attn = LLMAttention(weights_per_step=weights_per_step)
120 self.norm1 = RMSNorm()
121 self.norm2 = RMSNorm()
122 self.weights_per_step = weights_per_step
123 if self.weights_per_step:
124 self.gating = nn.ModuleList([ActivationGating(3072) for _ in range(9)])
125 else:
126 self.gating = ActivationGating()
127
128 def forward(self, x):
129 x = self.self_attn(self.norm1(x)) + x
130 if self.weights_per_step:
131 p = self.self_attn.k_history.shape[2] - 1
132 return x + self.gating[p if p < 9 else 8](self.norm2(x))
133 return x + self.gating(self.norm2(x))
134
135
136class LLMTransformer(nn.Module):
137
138 def __init__(
139 self,
140 num_layers=24,
141 weights_per_step=False):
142 super().__init__()
143 self.layers = nn.ModuleList(
144 [
145 LLMTransformerLayer(weights_per_step=weights_per_step)
146 for _ in range(num_layers)
147 ])
148
149 def forward(self, x):
150 for lay in self.layers:
151 x = lay(x)
152 return x
153
154
155class Voc(Wav2Vec2PreTrainedModel):
156
157 '''For using different batch_siz -> Voc._flush()
158 '''
159
160 def __init__(self, config=PretrainedConfig()):
161 super().__init__(config=config)
162 self.encoder_transformer = VocTransformer()
163 self.decoder_transformer = VocTransformer()
164 self.encoder = SEANetEncoder()
165 self.decoder = SEANetDecoder()
166 self.sample_rate = 24000
167 self.quantizer = SplitResidualVectorQuantizer()
168 self.downsample = BufferConv1d(512, 512, kernel_size=4, stride=2, groups=1, bias=False)
169 upsample_channel_wise_bug = True
170 self.upsample = BufferConvTranspose1d(512, 512, kernel_size=4,
171 groups=512 if upsample_channel_wise_bug else 1,
172 stride=2, bias=False)
173 self.frame_rate = 12.5
174 self.encode_buffer = None
175
176 def _flush(self):
177 '''stream buffers have tensors of old batch size! Voc()._flush() to clean buffers
178 '''
179 self.encode_buffer = None # holds unused (incomplete windows of len < 1920) - we need 1920 to produce 1 token
180 if self.downsample.previous is not None:
181 self.downsample.previous = None
182 if self.upsample.partial is not None:
183 self.upsample.partial = None
184 for arch in [self.encoder, self.decoder]:
185 for _m in arch.model:
186 if type(_m) is SEANetResnetBlock:
187 for _b in _m.block:
188 if type(_b) is BufferConv1d:
189 if _b.previous is not None:
190 _b.previous = None
191 if type(_m) is BufferConv1d:
192 if _m.previous is not None:
193 _m.previous = None
194 if type(_m) is BufferConvTranspose1d:
195 if _m.partial is not None:
196 _m.partial = None
197
198 @torch.no_grad()
199 def encode(self, x):
200 '''24KHz audio to codes
201 x : [bs, 1, 24 KHz]
202 c : [bs, 8, time] = 1920 audio samples produce 1 time frame (of n_q codebooks)
203 '''
204 if self.encode_buffer is not None:
205 x = torch.cat([self.encode_buffer, x], 2)
206 _bs, _1, _len = x.shape
207 num_frames = int(_len / 1920)
208 leftover = x[:, :, (num_frames+1) * 1920:]
209 if leftover.shape[2] > 0:
210 self.encode_buffer = leftover
211 else:
212 self.encode_buffer = None
213 torch.cuda.empty_cache()
214 if num_frames > 0:
215 c = []
216 for n in range(num_frames):
217 e = self.encoder(x[:, :, n * 1920:(n + 1) * 1920])
218 e = self.encoder_transformer(e)
219 e = self.downsample(e)
220 _c = self.quantizer.encode(e)
221 c.append(_c)
222 c = torch.cat(c, 2)
223 else:
224 # num_frames = 0 Early exit -> for x.shape[2]<1920 fill conv buffers but can't output token
225 c = torch.empty(_bs, 16, 0)
226 return c
227
228 @torch.no_grad()
229 def decode(self, c):
230 '''codes to 24kHZ audio
231 c: [bs, 8, n_tokens]
232 x: [bs, 1, n_tokens * 1920]
233 '''
234 _hidden = []
235 for i in range(c.shape[2]):
236 x = self.quantizer.decode(c[:, :, i:i+1])
237 x = self.upsample(x)
238 x = self.decoder_transformer(x)
239 x = self.decoder(x)
240 _hidden.append(x)
241 return torch.cat(_hidden, 2) # [bs, 1, 24KHz]
242
243
244class SEANetResnetBlock(nn.Module):
245 def __init__(
246 self,
247 dim,
248 kernel_sizes=[3, 1],
249 ):
250 super().__init__()
251
252 block = []
253 for i, kernel_size in enumerate(kernel_sizes):
254
255 block += [
256 nn.ELU(),
257 BufferConv1d(
258 dim if i == 0 else dim // 2,
259 dim // 2 if i == 0 else dim,
260 kernel_size=kernel_size,
261 bias=True,
262 ),
263 ]
264
265 self.block = nn.Sequential(*block)
266
267 def forward(self, x):
268 return x + self.block(x)
269
270
271class SEANetEncoder(nn.Module):
272 def __init__(
273 self,
274 channels=1, # DOES NOT SUPPORT STEREO
275 dimension=512,
276 n_filters=64,
277 ratios=[8, 6, 5, 4],
278 kernel_size=7,
279 last_kernel_size=3,
280 ):
281 super().__init__()
282 self.ratios = list(reversed(ratios))
283 del ratios
284 mult = 1
285 model=[
286 BufferConv1d(
287 channels,
288 mult * n_filters,
289 kernel_size,
290 bias=True
291 )
292 ]
293 for i, ratio in enumerate(self.ratios):
294 model += [SEANetResnetBlock(mult * n_filters),
295 nn.ELU(),
296 BufferConv1d(mult * n_filters,
297 mult * n_filters * 2,
298 kernel_size=ratio * 2,
299 stride=ratio,
300 bias=True)]
301 mult *= 2
302 # ENDFOR
303 model += [nn.ELU(),
304 BufferConv1d(mult * n_filters,
305 dimension,
306 last_kernel_size,
307 bias=True)]
308 self.model = nn.Sequential(*model)
309
310 def forward(self, x):
311 return self.model(x)
312
313
314class SEANetDecoder(nn.Module):
315
316 def __init__(
317 self,
318 channels=1,
319 dimension=512,
320 n_filters=64,
321 ratios=[8, 6, 5, 4],
322 kernel_size=7,
323 last_kernel_size=3):
324
325 super().__init__()
326 mult = int(2 ** len(ratios))
327 model = [BufferConv1d(dimension,
328 mult * n_filters,
329 kernel_size,
330 bias=True)]
331 #UP
332 for i, ratio in enumerate(ratios):
333 model += [nn.ELU(),
334 BufferConvTranspose1d(mult * n_filters,
335 mult * n_filters // 2,
336 kernel_size=ratio * 2,
337 stride=ratio,
338 bias=True),
339 SEANetResnetBlock(mult * n_filters // 2)]
340 mult //= 2
341 # LAST
342 model += [
343 nn.ELU(),
344 BufferConv1d(
345 n_filters,
346 channels,
347 last_kernel_size,
348 bias=True
349 ),
350 ]
351 self.model = nn.Sequential(*model)
352
353 def forward(self, x):
354 return self.model(x)
355
356
357class BufferConv1d(nn.Conv1d):
358 def __init__(self,
359 *args,
360 **kwargs):
361 super().__init__(*args, **kwargs)
362 self.previous = None
363
364 def forward(self, x):
365 k = self.kernel_size[0]
366
367 if self.previous is not None:
368
369 x = torch.cat([self.previous, x], 2)
370
371 else: # If self.previous is None => Use zero pad
372
373 if k == 3:
374
375 p = (2, 0)
376 x = F.pad(x, p, mode='replicate', value=0.0) # skip connections SeaNetResBlk
377
378 elif k == 4: # ConvTrUpsample is the first conv encountered by decode replicate solves pulse
379
380 p = (3, 0)
381 x = F.pad(x, p, mode='replicate', value=0.0)
382
383 elif k == 7:
384
385 p = (6, 0)
386 x = F.pad(x, p, mode='replicate', value=0.0)
387
388 elif k == 16:
389
390 p = (2, 0)
391 x = F.pad(x, p, mode='replicate', value=0.0) # THis can be also constant w/o pulse occur
392
393 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
394 offset = num_frames * self.stride[0]
395 self.previous = x[..., offset:]
396 return super().forward(x)
397
398
399class BufferConvTranspose1d(nn.ConvTranspose1d):
400 # kernel 5 has only 1 pixel for input (cloned)
401 # https://distill.pub/2016/deconv-checkerboard/
402 def __init__(self,
403 *args,
404 **kwargs):
405 super().__init__(*args,
406 **kwargs)
407 self.partial = None
408
409 def forward(self, x):
410 out = super().forward(x)
411 OT = out.shape[2]
412 invalid_steps = self.kernel_size[0] - self.stride[0]
413 if self.partial is not None:
414 PT = self.partial.shape[-1]
415 if self.bias is not None:
416 out[..., :PT] += self.partial - self.bias[:, None]
417 else:
418 out[..., :PT] += self.partial # for ConvTrUpsample1d
419 invalid_steps = self.kernel_size[0] - self.stride[0]
420 self.partial = out[..., OT - invalid_steps :]
421 out = out[...,:OT - invalid_steps]
422 return out
423
424
425class CodeBook(nn.Module):
426 def __init__(self, dim, codebook_size):
427 super().__init__()
428 self.register_buffer('_e', torch.zeros(codebook_size, dim))
429
430 def encode(self, x):
431 dist = torch.cdist(
432 x.transpose(1, 2), # [bs, time, 256]
433 self._e[None, :, :] # [1, 2048, 256]
434 )
435 codes = dist.argmin(2)
436 return codes
437
438 def decode(self, codes):
439 quantized = F.embedding(codes, self._e)
440 return quantized.transpose(1, 2) # [1, 256, time]
441
442
443class SplitResidualVectorQuantizer(nn.Module):
444
445 def __init__(self,
446 n_q=None):
447 super().__init__()
448 self.in_proj_s = torch.nn.Conv1d(512, 256, 1, bias=False)
449 self.in_proj_a = torch.nn.Conv1d(512, 256, 1, bias=False)
450 self.out_proj_s = torch.nn.Conv1d(256, 512, 1, bias=False) # reused for all _acoustic_books
451 self.out_proj_a = torch.nn.Conv1d(256, 512, 1, bias=False)
452 self.layers = nn.ModuleList([CodeBook(dim=256, codebook_size=2048) for _ in range(18)])
453 self._acoustic_books = range(1, 16) # Official Mimi
454 # CODEBOOKS
455 # Here we re use RVQ codebooks for higher fidelity!
456 # Exclude 0 here as it has different proj (in_proj_s)
457 # self._acoustic_books = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 17, 17, 17, 17]
458
459 def encode(self, x):
460 indices = self.layers[0].encode(self.in_proj_s(x)) # integers
461 all_indices = [ indices[:, None, :], ]
462 x = self.in_proj_a(x)
463 for _cb in self._acoustic_books:
464 indices = self.layers[_cb].encode(x)
465 x = x - self.layers[_cb].decode(indices)
466 all_indices.append(indices[:, None, :])
467 codes = torch.cat(all_indices, 1)
468 return codes
469
470 def decode(self, codes):
471 _s = self.layers[0].decode(codes[:, 0, :])
472 _a = torch.zeros([1, 1], device=codes.device)
473 for i, _cb in enumerate(self._acoustic_books):
474 _a = _a + self.layers[_cb].decode(codes[:, i+1, :])
475 return self.out_proj_s(_s) + self.out_proj_a(_a) # [bs, 512, time]
476
477
478class VocAttention(nn.Module):
479
480 def __init__(self,
481 embed_dim):
482
483 super().__init__()
484 self.fused_proj = nn.Parameter(torch.zeros(embed_dim, embed_dim))
485
486 def forward(self, x):
487 '''bypass of streaming training'''
488 if x.shape[1] > 1:
489 x = x.mean(1, keepdims=True)
490 x = torch.matmul(x, self.fused_proj)
491 return x # FFN broadcasts to x.shape[1]=2
492
493
494class VocTransformerLayer(nn.Module):
495
496 def __init__(self, d_model=512, dim_feedforward=2048):
497 super().__init__()
498 self.self_attn = VocAttention(embed_dim=d_model)
499 self.norm1 = nn.LayerNorm(d_model, eps=1e-5)
500 self.norm2 = nn.LayerNorm(d_model, eps=1e-5)
501 self.linear1 = nn.Linear(d_model, dim_feedforward, bias=False)
502 self.linear2 = nn.Linear(dim_feedforward, d_model, bias=False)
503
504 def forward(self, x):
505 x = x + self.self_attn(self.norm1(x))
506 return x + self.linear2(F.gelu(self.linear1(self.norm2(x))))
507
508
509class VocTransformer(nn.Module):
510
511 def __init__(self):
512
513 super().__init__()
514 self.layers = nn.ModuleList(VocTransformerLayer() for _ in range(8))
515
516 def forward(self, x):
517 x = x.transpose(1, 2)
518 for la in self.layers:
519 x = la(x)
520 return x.transpose(1, 2)
521
522class Entry():
523 def __init__(self, tokens=None):
524 self.tokens = tokens
525 self.padding = len(tokens) + 2 - 1
526
527class TokenState:
528
529 def __init__(self, entries = None):
530 self.entries = entries
531 self.queued = deque([])
532 self.lookahead_queued = deque()
533 self.end_step = None
534 self.forced_padding = 2
535
536class TTSModel(nn.Module):
537
538 def __init__(self):
539 super().__init__()
540 self.tokenizer = SentencePieceProcessor(str(hf_hub_download(repo_id='kyutai/tts-0.75b-en-public',
541 filename='tokenizer_spm_8k_en_fr_audio.model')))
542 with torch.device("meta"):
543 self.emb = nn.ModuleList([ScaledEmbedding(2049, 1024) for _ in range(16)])
544 self.text_emb = ScaledEmbedding(8001, 1024, demux_second_stream=True)
545 self.transformer = LLMTransformer()
546 self.out_norm = RMSNorm()
547 self.depformer_in = nn.ModuleList([nn.Linear(1024, 1024, bias=False) for _ in range(9)])
548 self.depformer_emb = nn.ModuleList([ScaledEmbedding(2049, 128) for _ in range(16 - 1)])
549 self.depformer_text_emb = ScaledEmbedding(8001, 128, demux_second_stream=True)
550 self.depformer = LLMTransformer(num_layers=4, weights_per_step=16)
551 self.linears = nn.ModuleList([nn.Linear(1024, 2048, bias=False) for _ in range(16)]) # DPF heads
552
553 state_d = load_file(hf_hub_download(repo_id='Dionyssos/_TTS075B', filename='tts_075B.safetensors'))
554 self.load_state_dict(state_d, assign=True, strict=True) #overwrite devices of rand init params
555 self.to(dtype=torch.bfloat16).eval()
556
557 def prepare_script(self, script='Type your text here.'):
558 entries = []
559 # break is indicated as e.g. <break time="3s"/>
560 event_re = re.compile(r"(?:<break\s+time=\"([0-9]+(?:.[0-9]*)?)s\"\s*/?>)|(?:\s+)")
561 line = script.replace('’', "'").replace(':', " ").replace('(', "").replace(')', "")
562 while line:
563 match = event_re.search(line)
564 if match is None:
565 break
566 word = line[:match.start()]
567 line = line[match.end():]
568 if word:
569 entries.append(Entry(tokens=self.tokenizer.encode(word)))
570 if match.group(1):
571 raise ValueError
572 # break_duration = float(match.group(1))
573 # padding = int(round(break_duration * frame_rate))
574 # entry = Entry(tokens=[], text='', padding=padding)
575 # entries.append(entry)
576 if line:
577 entries.append(Entry(tokens=self.tokenizer.encode(line)))
578 return entries
579
580 @property
581 def device(self):
582 return next(iter(self.parameters())).device
583
584 @torch.no_grad()
585 def generate(self, text=None,
586 voice_path=None, mimi=None,
587 play=16):
588 _wav, _ = sphn.read(voice_path,
589 sample_rate=24000)
590 _wav = mimi.encode(torch.from_numpy(_wav).to(device=self.device)[None])[0, :, :] # limit frames of voice prefix
591 state = TokenState(entries=deque(self.prepare_script(script=text)))
592 upper_lim = 2 * sum([len(p.tokens) for p in state.entries])
593 self.cache = torch.full((2,17, 4), -1, device=self.device, dtype=torch.long)
594 pcms = []#final audio to return
595 for offset in range(4 * upper_lim):
596 print(f'{offset=} of {upper_lim=}',end='\r')
597 if state.end_step is not None:
598 if offset >= state.end_step + 16 + 4:
599 break
600
601 input_ = self.cache[:, :, offset % self.cache.shape[2]].clone()
602
603 if offset == 0:
604 input_[:, 0] = 8000 # so we dont have to reset cfg txr = -1 for offset >0
605 input_[:, 1:] = 2048
606
607 if offset < 3:
608 input_[:, 2:] = 2048
609
610
611 x = self.text_emb(input_[:, :1])
612 for cb_ in range(16):
613 x = self.emb[cb_](input_[:, cb_ + 1 : cb_ + 2]) + x
614 x = self.out_norm(self.transformer(x))
615
616
617 token = -1
618 if offset > _wav.shape[1]:
619 token = 0
620 # START
621 if state.queued:
622 token = 3
623 if state.forced_padding > 0:
624 token = 3
625 #===================================
626 if token == 0:
627 if state.entries:
628 e = state.entries.popleft()
629 if e.tokens:
630 state.queued.extend(e.tokens)
631 lookahead =2
632 for e2 in state.entries:
633 if e2.tokens:
634 lookahead -= 1
635 if lookahead == 0:
636 state.lookahead_queued.extend(e2.tokens)
637 break
638 # print('\neeee',e2,'\n\n')
639 # raise ValueError
640 else:
641 token = 3
642 state.forced_padding = e.padding
643 # print(f'\n\n=========o=============\n{state.lookahead_queued=} {state.queued=}===================\n\n')
644 else:
645 token = 3
646 if state.end_step is None:
647 token = 0
648 if state.end_step is None:
649 state.end_step = offset
650 #==============================================
651 output=0
652 if token == 3:
653 if state.forced_padding > 0:
654 state.forced_padding -= 1
655 if state.queued:
656 output = state.queued.popleft()
657 else:
658 output = 3
659 # ==========================
660 second = -1
661 if output == 0:
662 second = 0
663 if state.queued:
664 output = state.queued.popleft()
665 else:
666 output = 3
667 elif state.lookahead_queued:
668 second = state.lookahead_queued.popleft() # Difference of queued and lookahead_queued?
669 token = (second + 1) * 8001 + output
670
671 # audio tokens
672 ac = (offset + 1) % self.cache.shape[2]
673 self.cache[0, 0, ac] = token
674 audio_tokens = torch.ones([1, 16], device=x.device, dtype=torch.long)
675 if offset > play:
676 prev_token = torch.tensor([[token]], device=x.device, dtype=torch.long)
677 for _cb in range(16):
678 last_token_input = None
679 if _cb == 0:
680 last_token_input = self.depformer_text_emb(prev_token.repeat(2, 1))
681 else:
682 last_token_input = self.depformer_emb[_cb - 1](prev_token)
683 dep_output = self.depformer(self.depformer_in[_cb if _cb < 9 else 8](x) + last_token_input)
684 logits = self.linears[_cb](dep_output)
685 prev_token = (2.0 * logits[0, :, :] - logits[1, :, :]).argmax(1)
686 audio_tokens[0, _cb] = prev_token
687 # voXcopy
688 if offset > play and offset < play + 1 + _wav.shape[1]:
689 audio_tokens[:, 0] = _wav[0, offset - play - 1]
690 if offset > play and offset < play + 2 + _wav.shape[1]:
691 audio_tokens[:, 1:] = _wav[1:, offset - play - 2]
692 # next turn
693 self.cache[0, 1:, ac] = audio_tokens
694 # cfg
695 if offset > 16 + 2 + _wav.shape[1]:
696 if offset > 16 + 4 + _wav.shape[1]:
697 self.cache[1, 1:, ac] = self.cache[0, 1:, ac]
698 else:
699 self.cache[1, 1, ac] = self.cache[0, 1, ac]
700 # ivao0/voc
701 if offset > 20 + _wav.shape[1]:
702 audio_tokens[:, 0] = self.cache[0, 1, (offset - 1) % self.cache.shape[2]] # previous
703 pcms.append(mimi.decode(audio_tokens[:, :, None])) # [1,1,1920]
704 x = torch.cat(pcms, dim=2)[0, 0, :]
705 return x.cpu().numpy()
706
707class ScaledEmbedding(nn.Embedding):
708 def __init__(self, num_embeddings=None, embedding_dim=None, demux_second_stream=False):
709 super().__init__(num_embeddings, embedding_dim)
710 self.zero_idx = -1
711 self.low_rank = None
712 self.demux_second_stream = demux_second_stream
713 if self.demux_second_stream:
714 self.out1 = nn.Linear(embedding_dim, 1024, bias=False)
715 self.out2 = nn.Linear(embedding_dim, 1024, bias=False)
716 else:
717 if embedding_dim != 1024:
718 self.low_rank = nn.Linear(embedding_dim, 1024, bias=False)
719
720 def forward(self, input):
721 is_zero = input == self.zero_idx
722 zero = torch.zeros(1, dtype=input.dtype, device=input.device)
723 input = input.clamp(min=0)
724 if self.demux_second_stream:
725 left = super().forward(input % self.num_embeddings)
726 right = input // self.num_embeddings - 1
727 right_zero = (right < 0)[..., None]
728 right.clamp_(min=0)
729 right = super().forward(right)
730 y = self.out1(left) + torch.where(right_zero, zero, self.out2(right))
731 y = torch.where(is_zero[..., None], zero, y)
732 else:
733 y = super().forward(input)
734 y = torch.where(is_zero[..., None], zero, y)
735 if self.low_rank is not None:
736 # Can only see low_rank if no demux second stream
737 y = self.low_rank(y) # applies after
738 return y
739
740text = '''Far over the misty mountains cold
741To dungeons deep and caverns old
742We must away ere break of day
743To seek the pale enchanted gold.
744
745The dwarves of yore made mighty spells,
746While hammers fell like ringing bells
747In places deep, where dark things sleep,
748In hollow halls beneath the fells.
749
750For ancient king and elvish lord
751There many a gleaming golden hoard
752They shaped and wrought, and light they caught
753To hide in gems on hilt of sword.
754
755On silver necklaces they strung
756The flowering stars, on crowns they hung
757The dragon-fire, in twisted wire
758They meshed the light of moon and sun.
759
760Far over the misty mountains cold
761To dungeons deep and caverns old
762We must away, ere break of day,
763To claim our long-forgotten gold.
764Farewell we call to hearth and hall!
765Though wind may blow and rain may fall,
766We must away ere break of day
767Far over wood and mountain tall.'''
768
769
770device = 'cpu' # 'cuda:0'
771tts_model = TTSModel().eval().to(device)
772mimi = Voc.from_pretrained('ivao0/voc').eval().to(device)
773x = tts_model.generate(text=text,
774 voice_path=hf_hub_download(repo_id='Dionyssos/_TTS075B', filename='wav/en_US_m-ailabs_mary_ann.wav'),
775 mimi=mimi)
776sphn.write_wav(f'dsm_tts.wav', x, 24000)
777