deviandice/DL-Art-School
1
0
Fork 0
forked from mrq/DL-Art-School
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Add code validation to autoregressive_codegen

Browse Source
This commit is contained in:
James Betker 2022-04-04 09:51:41 -06:00
parent 99de63a922
commit cdd12ff46c
1 changed files with 5 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

5
codes/models/audio/tts/autoregressive_codegen.py
View File

@ -201,6 +201,8 @@ class AutoregressiveCodegen(nn.Module):
self.START_TOKEN=8192
self.STOP_TOKEN=8193
self.max_text_token_id = num_text_tokens
self.max_mel_token_id = num_mel_tokens
self.mel_embedding = ConditioningEncoder(80, model_dim, do_checkpointing=False)
self.encoder = CheckpointedXTransformerWrapper(
num_tokens=num_text_tokens,
@ -243,6 +245,9 @@ class AutoregressiveCodegen(nn.Module):
}
def forward(self, text_codes, conditioning_signal, mel_codes, wav_lengths, return_loss=True):
assert text_codes.max() < self.max_text_token_id and text_codes.min() >= 0, f'Invalid text code encountered: {text_codes.max()}, {text_codes.min()}'
assert mel_codes.max() < self.max_mel_token_id and mel_codes.min() >= 0, f'Invalid mel code encountered: {mel_codes.max()}, {mel_codes.min()}'
# Format mel_codes with a stop token on the end.
mel_lengths = wav_lengths // 1024 + 1
for b in range(mel_codes.shape[0]):