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

nevermind thats slow

Browse Source
This commit is contained in:
mrq 2025-02-14 16:35:17 -06:00
parent 285e493b12
commit 13c3a08853
1 changed files with 2 additions and 9 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

11
vall_e/models/base.py
View File

@ -370,16 +370,12 @@ class AudioEncoder(nn.Module):
class AudioDecoder(nn.Module): class AudioDecoder(nn.Module):
def __init__( def __init__(
self, self,
levels,
d_model, d_model,
hidden_size, hidden_size,
vocab_size, vocab_size,
): ):
super().__init__() super().__init__()
hidden_size *= levels
vocab_size *= levels
self.vocab_size = vocab_size self.vocab_size = vocab_size
self.up = nn.Linear( d_model, hidden_size ) self.up = nn.Linear( d_model, hidden_size )
self.down = nn.Linear( hidden_size, vocab_size ) self.down = nn.Linear( hidden_size, vocab_size )
@ -715,8 +711,6 @@ class Base(nn.Module):
self.resp_parallel_training = True # governs if all levels are trained in parallel or one per sample like the old way self.resp_parallel_training = True # governs if all levels are trained in parallel or one per sample like the old way
self.monolithic_audio_encoder = False # monolithic sounds bad self.monolithic_audio_encoder = False # monolithic sounds bad
if self.version >= 7: if self.version >= 7:
dec_dim = d_model * 4
if self.monolithic_audio_encoder: if self.monolithic_audio_encoder:
self.audio_emb = AudioEncoder( self.audio_emb = AudioEncoder(
n_tokens=n_audio_tokens + 1, # masked token n_tokens=n_audio_tokens + 1, # masked token
@ -736,10 +730,9 @@ class Base(nn.Module):
) )
self.audio_decoder = AudioDecoder( self.audio_decoder = AudioDecoder(
self.n_resp_levels,
d_model, d_model,
dec_dim, d_model * 2,
n_audio_tokens + 1, (n_audio_tokens + 1) * self.n_resp_levels,
) )
if attention_backend == "auto": if attention_backend == "auto":