Spaces:
Sleeping
Sleeping
# modified from https://github.com/feng-yufei/shared_debugging_code/blob/main/model/t2s_model.py | |
import torch | |
from tqdm import tqdm | |
from AR.models.utils import make_pad_mask | |
from AR.models.utils import topk_sampling,sample,logits_to_probs,multinomial_sample_one_no_sync | |
from AR.modules.embedding import SinePositionalEmbedding | |
from AR.modules.embedding import TokenEmbedding | |
from AR.modules.transformer import LayerNorm | |
from AR.modules.transformer import TransformerEncoder | |
from AR.modules.transformer import TransformerEncoderLayer | |
from torch import nn | |
from torch.nn import functional as F | |
from torchmetrics.classification import MulticlassAccuracy | |
default_config = { | |
"embedding_dim": 512, | |
"hidden_dim": 512, | |
"num_head": 8, | |
"num_layers": 12, | |
"num_codebook": 8, | |
"p_dropout": 0.0, | |
"vocab_size": 1024 + 1, | |
"phoneme_vocab_size": 512, | |
"EOS": 1024 | |
} | |
class Text2SemanticDecoder(nn.Module): | |
def __init__(self, config, norm_first=False, top_k=3): | |
super(Text2SemanticDecoder, self).__init__() | |
self.model_dim = config['model']["hidden_dim"] | |
self.embedding_dim = config['model']["embedding_dim"] | |
self.num_head = config['model']["head"] | |
self.num_layers = config['model']["n_layer"] | |
self.norm_first = norm_first | |
self.vocab_size = config['model']["vocab_size"] | |
self.phoneme_vocab_size = config['model']["phoneme_vocab_size"] | |
self.p_dropout = config['model']["dropout"] | |
self.EOS = config['model']["EOS"] | |
self.norm_first = norm_first | |
assert self.EOS == self.vocab_size - 1 | |
# should be same as num of kmeans bin | |
# assert self.EOS == 1024 | |
self.bert_proj = nn.Linear(1024, self.embedding_dim) | |
self.ar_text_embedding = TokenEmbedding( | |
self.embedding_dim, self.phoneme_vocab_size, self.p_dropout) | |
self.ar_text_position = SinePositionalEmbedding( | |
self.embedding_dim, dropout=0.1, scale=False, alpha=True) | |
self.ar_audio_embedding = TokenEmbedding( | |
self.embedding_dim, self.vocab_size, self.p_dropout) | |
self.ar_audio_position = SinePositionalEmbedding( | |
self.embedding_dim, dropout=0.1, scale=False, alpha=True) | |
self.h = TransformerEncoder( | |
TransformerEncoderLayer( | |
d_model=self.model_dim, | |
nhead=self.num_head, | |
dim_feedforward=self.model_dim * 4, | |
dropout=0.1, | |
batch_first=True, | |
norm_first=norm_first, ), | |
num_layers=self.num_layers, | |
norm=LayerNorm(self.model_dim) if norm_first else None, ) | |
self.ar_predict_layer = nn.Linear( | |
self.model_dim, self.vocab_size, bias=False) | |
self.loss_fct = nn.CrossEntropyLoss(reduction='sum') | |
self.ar_accuracy_metric = MulticlassAccuracy( | |
self.vocab_size, | |
top_k=top_k, | |
average="micro", | |
multidim_average="global", | |
ignore_index=self.EOS, ) | |
def forward(self, x, x_lens, y, y_lens, bert_feature): | |
''' | |
x: phoneme_ids | |
y: semantic_ids | |
''' | |
x = self.ar_text_embedding(x) | |
x = x + self.bert_proj(bert_feature.transpose(1,2)) | |
x = self.ar_text_position(x) | |
x_mask = make_pad_mask(x_lens) | |
y_mask = make_pad_mask(y_lens) | |
y_mask_int = y_mask.type(torch.int64) | |
codes = y.type(torch.int64) * (1 - y_mask_int) | |
# Training | |
# AR Decoder | |
y, targets = self.pad_y_eos(codes, y_mask_int, eos_id=self.EOS) | |
x_len = x_lens.max() | |
y_len = y_lens.max() | |
y_emb = self.ar_audio_embedding(y) | |
y_pos = self.ar_audio_position(y_emb) | |
xy_padding_mask = torch.concat([x_mask, y_mask], dim=1) | |
ar_xy_padding_mask = xy_padding_mask | |
x_attn_mask = F.pad( | |
torch.zeros((x_len, x_len), dtype=torch.bool, device=x.device), | |
(0, y_len), | |
value=True, ) | |
y_attn_mask = F.pad( | |
torch.triu( | |
torch.ones(y_len, y_len, dtype=torch.bool, device=x.device), | |
diagonal=1, ), | |
(x_len, 0), | |
value=False, ) | |
xy_attn_mask = torch.concat([x_attn_mask, y_attn_mask], dim=0) | |
bsz, src_len = x.shape[0], x_len + y_len | |
_xy_padding_mask = (ar_xy_padding_mask.view(bsz, 1, 1, src_len) | |
.expand(-1, self.num_head, -1, -1) | |
.reshape(bsz * self.num_head, 1, src_len)) | |
xy_attn_mask = xy_attn_mask.logical_or(_xy_padding_mask) | |
new_attn_mask = torch.zeros_like(xy_attn_mask, dtype=x.dtype) | |
new_attn_mask.masked_fill_(xy_attn_mask, float("-inf")) | |
xy_attn_mask = new_attn_mask | |
# x 和完整的 y 一次性输入模型 | |
xy_pos = torch.concat([x, y_pos], dim=1) | |
xy_dec, _ = self.h( | |
(xy_pos, None), | |
mask=xy_attn_mask, ) | |
logits = self.ar_predict_layer(xy_dec[:, x_len:]).permute(0, 2, 1) | |
# loss | |
# from feiteng: 每次 duration 越多, 梯度更新也应该更多, 所以用 sum | |
loss = F.cross_entropy(logits, targets, reduction='sum') | |
acc = self.ar_accuracy_metric(logits.detach(), targets).item() | |
return loss, acc | |
# 需要看下这个函数和 forward 的区别以及没有 semantic 的时候 prompts 输入什么 | |
def infer(self, | |
x, | |
x_lens, | |
prompts, | |
bert_feature, | |
top_k: int=-100, | |
early_stop_num: int=-1, | |
temperature: float=1.0): | |
x = self.ar_text_embedding(x) | |
x = x + self.bert_proj(bert_feature.transpose(1,2)) | |
x = self.ar_text_position(x) | |
# AR Decoder | |
y = prompts | |
prefix_len = y.shape[1] | |
x_len = x.shape[1] | |
x_attn_mask = torch.zeros((x_len, x_len), dtype=torch.bool) | |
stop = False | |
for _ in tqdm(range(1500)): | |
y_emb = self.ar_audio_embedding(y) | |
y_pos = self.ar_audio_position(y_emb) | |
# x 和逐渐增长的 y 一起输入给模型 | |
xy_pos = torch.concat([x, y_pos], dim=1) | |
y_len = y.shape[1] | |
x_attn_mask_pad = F.pad( | |
x_attn_mask, | |
(0, y_len), | |
value=True, ) | |
y_attn_mask = F.pad( | |
torch.triu( | |
torch.ones(y_len, y_len, dtype=torch.bool), diagonal=1), | |
(x_len, 0), | |
value=False, ) | |
xy_attn_mask = torch.concat( | |
[x_attn_mask_pad, y_attn_mask], dim=0).to(y.device) | |
xy_dec, _ = self.h( | |
(xy_pos, None), | |
mask=xy_attn_mask, ) | |
logits = self.ar_predict_layer(xy_dec[:, -1]) | |
samples = topk_sampling( | |
logits, top_k=top_k, top_p=1.0, temperature=temperature) | |
if early_stop_num != -1 and (y.shape[1] - prefix_len | |
) > early_stop_num: | |
print("use early stop num:", early_stop_num) | |
stop = True | |
if torch.argmax( | |
logits, dim=-1)[0] == self.EOS or samples[0, 0] == self.EOS: | |
# print(torch.argmax(logits, dim=-1)[0] == self.EOS, samples[0, 0] == self.EOS) | |
stop = True | |
if stop: | |
if prompts.shape[1] == y.shape[1]: | |
y = torch.concat([y, torch.zeros_like(samples)], dim=1) | |
print('bad zero prediction') | |
print(f"T2S Decoding EOS [{prefix_len} -> {y.shape[1]}]") | |
break | |
# 本次生成的 semantic_ids 和之前的 y 构成新的 y | |
# print(samples.shape)#[1,1]#第一个1是bs | |
# import os | |
# os._exit(2333) | |
y = torch.concat([y, samples], dim=1) | |
return y | |
def pad_y_eos(self, y, y_mask_int, eos_id): | |
targets = F.pad( | |
y, (0, 1), value=0) + eos_id * F.pad( | |
y_mask_int, (0, 1), value=1) | |
# 错位 | |
return targets[:, :-1], targets[:, 1:] | |
def infer_panel(self, | |
x,#####全部文本token | |
x_lens, | |
prompts,####参考音频token | |
bert_feature, | |
top_k: int=-100, | |
early_stop_num: int=-1, | |
temperature: float=1.0): | |
x = self.ar_text_embedding(x) | |
x = x + self.bert_proj(bert_feature.transpose(1,2)) | |
x = self.ar_text_position(x) | |
# AR Decoder | |
y = prompts | |
prefix_len = y.shape[1] | |
x_len = x.shape[1] | |
x_attn_mask = torch.zeros((x_len, x_len), dtype=torch.bool) | |
stop = False | |
# print(1111111,self.num_layers) | |
cache={ | |
"all_stage":self.num_layers, | |
"k":[None]*self.num_layers,###根据配置自己手写 | |
"v":[None]*self.num_layers, | |
# "xy_pos":None,##y_pos位置编码每次都不一样的没法缓存,每次都要重新拼xy_pos.主要还是写法原因,其实是可以历史统一一样的,但也没啥计算量就不管了 | |
"y_emb":None,##只需要对最新的samples求emb,再拼历史的就行 | |
# "logits":None,###原版就已经只对结尾求再拼接了,不用管 | |
# "xy_dec":None,###不需要,本来只需要最后一个做logits | |
"first_infer":1, | |
"stage":0 | |
} | |
for idx in tqdm(range(1500)): | |
if(cache["first_infer"]==1): | |
y_emb = self.ar_audio_embedding(y) | |
else: | |
y_emb = torch.cat([cache["y_emb"],self.ar_audio_embedding(y[:,-1:])],1) | |
cache["y_emb"]=y_emb | |
y_pos = self.ar_audio_position(y_emb) | |
# x 和逐渐增长的 y 一起输入给模型 | |
if(cache["first_infer"]==1): | |
xy_pos = torch.concat([x, y_pos], dim=1) | |
else: | |
xy_pos=y_pos[:,-1:] | |
y_len = y_pos.shape[1] | |
###以下3个不做缓存 | |
if (cache["first_infer"] == 1): | |
x_attn_mask_pad = F.pad( | |
x_attn_mask, | |
(0, y_len),###xx的纯0扩展到xx纯0+xy纯1,(x,x+y) | |
value=True, ) | |
y_attn_mask = F.pad(###yy的右上1扩展到左边xy的0,(y,x+y) | |
torch.triu( | |
torch.ones(y_len, y_len, dtype=torch.bool), diagonal=1), | |
(x_len, 0), | |
value=False, ) | |
xy_attn_mask = torch.concat( | |
[x_attn_mask_pad, y_attn_mask], dim=0).to(y.device) | |
else: | |
###最右边一列(是错的) | |
# xy_attn_mask=torch.ones((1, x_len+y_len), dtype=torch.bool,device=xy_pos.device) | |
# xy_attn_mask[:,-1]=False | |
###最下面一行(是对的) | |
xy_attn_mask = torch.zeros((1, x_len + y_len), dtype=torch.bool, device=xy_pos.device) | |
# pdb.set_trace() | |
###缓存重头戏 | |
# print(1111,xy_pos.shape,xy_attn_mask.shape,x_len,y_len) | |
xy_dec, _ = self.h( | |
(xy_pos, None), | |
mask=xy_attn_mask,cache=cache ) | |
logits = self.ar_predict_layer(xy_dec[:, -1])##不用改,如果用了cache的默认就是只有一帧,取最后一帧一样的 | |
# samples = topk_sampling(logits, top_k=top_k, top_p=1.0, temperature=temperature) | |
samples = sample(logits[0], y, top_k=top_k, top_p=1.0, repetition_penalty=1.35)[0].unsqueeze(0) | |
if early_stop_num != -1 and (y.shape[1] - prefix_len | |
) > early_stop_num: | |
print("use early stop num:", early_stop_num) | |
stop = True | |
if torch.argmax( | |
logits, dim=-1)[0] == self.EOS or samples[0, 0] == self.EOS: | |
# print(torch.argmax(logits, dim=-1)[0] == self.EOS, samples[0, 0] == self.EOS) | |
stop = True | |
if stop: | |
if prompts.shape[1] == y.shape[1]: | |
y = torch.concat([y, torch.zeros_like(samples)], dim=1) | |
print('bad zero prediction') | |
print(f"T2S Decoding EOS [{prefix_len} -> {y.shape[1]}]") | |
break | |
# 本次生成的 semantic_ids 和之前的 y 构成新的 y | |
# print(samples.shape)#[1,1]#第一个1是bs | |
y = torch.concat([y, samples], dim=1) | |
cache["first_infer"]=0 | |
return y,idx | |