diff --git a/style_bert_vits2/models/infer.py b/style_bert_vits2/models/infer.py index 835a042..c07b8c0 100644 --- a/style_bert_vits2/models/infer.py +++ b/style_bert_vits2/models/infer.py @@ -224,14 +224,23 @@ def infer( en_bert = en_bert[:, :-2] with torch.no_grad(): + # BERT may run in fp16 (for example on Colab) while a loaded TTS model + # remains fp32. Conv1d/Linear require inputs and parameters to use the + # same dtype outside autocast, so normalize all floating conditioning + # tensors to the TTS model dtype at this boundary. + model_dtype = next(net_g.parameters()).dtype x_tst = phones.to(device).unsqueeze(0) tones = tones.to(device).unsqueeze(0) lang_ids = lang_ids.to(device).unsqueeze(0) - bert = bert.to(device).unsqueeze(0) - ja_bert = ja_bert.to(device).unsqueeze(0) - en_bert = en_bert.to(device).unsqueeze(0) + bert = bert.to(device=device, dtype=model_dtype).unsqueeze(0) + ja_bert = ja_bert.to(device=device, dtype=model_dtype).unsqueeze(0) + en_bert = en_bert.to(device=device, dtype=model_dtype).unsqueeze(0) x_tst_lengths = torch.LongTensor([phones.size(0)]).to(device) - style_vec_tensor = torch.from_numpy(style_vec).to(device).unsqueeze(0) + style_vec_tensor = ( + torch.from_numpy(style_vec) + .to(device=device, dtype=model_dtype) + .unsqueeze(0) + ) del phones sid_tensor = torch.LongTensor([sid]).to(device)