mrq/vall-e
5
9
Fork 0
Code Issues 8 Pull Requests Packages Projects Releases Wiki Activity

uh

Browse Source
This commit is contained in:
mrq 2024-05-27 19:05:56 -05:00
parent 197d517181
commit c0ac84c795
1 changed files with 9 additions and 17 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

26
vall_e/models/base.py
View File

@ -905,9 +905,9 @@ class Base(nn.Module):
quant_level = quant_levels[bi] if quant_levels is not None else None quant_level = quant_levels[bi] if quant_levels is not None else None
if name in ["text" ]: if name == "text":
text_batch.append( input ) text_batch.append( input )
elif name == "prom": # and (quant_level is None or quant_level == 0) and not self.config.audio_embedding_sums: elif name == "prom":
prom_batch.append( input[:, quant_level] if quant_level is not None else input ) prom_batch.append( input[:, quant_level] if quant_level is not None else input )
elif name == "targ": elif name == "targ":
resp_batch.append( input ) resp_batch.append( input )
@ -1001,23 +1001,15 @@ class Base(nn.Module):
logits_prom = [] logits_prom = []
logits_resp = [] logits_resp = []
# trim logits to each section
for i, logit in enumerate(logits): for i, logit in enumerate(logits):
text_len = target_text_list[i].shape[0] if target_text_list else 0 text_len = target_text_list[i].shape[0]
prom_len = target_prom_list[i].shape[0] if target_prom_list else 0 prom_len = target_prom_list[i].shape[0]
resp_len = target_resp_list[i].shape[0] if target_resp_list else 0 resp_len = target_resp_list[i].shape[0]
if target_text_list: logits_text.append( logit[:text_len] )
logit_text = logit[:text_len] logits_prom.append( logit[text_len+1:text_len+1+prom_len] ) # + 1 to include separator
logits_text.append( logit_text ) logits_resp.append( logit[-resp_len:] )
# + 1 to include separator
if target_prom_list:
logit_prom = logit[text_len+1:text_len+1+prom_len]
logits_prom.append( logit_prom )
if target_resp_list:
logit_resp = logit[-resp_len:]
logits_resp.append( logit_resp )
# modify only for the AR so it can properly behave like a transformer # modify only for the AR so it can properly behave like a transformer