diff --git a/TTS/tts/layers/bark/hubert/kmeans_hubert.py b/TTS/tts/layers/bark/hubert/kmeans_hubert.py index c7724c23..a6a3b9ae 100644 --- a/TTS/tts/layers/bark/hubert/kmeans_hubert.py +++ b/TTS/tts/layers/bark/hubert/kmeans_hubert.py @@ -15,6 +15,8 @@ from einops import pack, unpack from torch import nn from torchaudio.functional import resample from transformers import HubertModel + + def round_down_nearest_multiple(num, divisor): return num // divisor * divisor diff --git a/TTS/tts/models/vits.py b/TTS/tts/models/vits.py index bc96f5dc..f4f4c639 100644 --- a/TTS/tts/models/vits.py +++ b/TTS/tts/models/vits.py @@ -1875,7 +1875,11 @@ class Vits(BaseTTS): def load_onnx(self, model_path: str, cuda=False): import onnxruntime as ort - providers = ["CPUExecutionProvider" if cuda is False else "CUDAExecutionProvider"] + providers = [ + "CPUExecutionProvider" + if cuda is False + else ("CUDAExecutionProvider", {"cudnn_conv_algo_search": "DEFAULT"}) + ] sess_options = ort.SessionOptions() self.onnx_sess = ort.InferenceSession( model_path, @@ -1883,11 +1887,8 @@ class Vits(BaseTTS): providers=providers, ) - def inference_onnx(self, x, x_lengths=None): - """ONNX inference (only single speaker models are supported) - - TODO: implement multi speaker support. - """ + def inference_onnx(self, x, x_lengths=None, speaker_id=None): + """ONNX inference""" if isinstance(x, torch.Tensor): x = x.cpu().numpy() @@ -1907,7 +1908,7 @@ class Vits(BaseTTS): "input": x, "input_lengths": x_lengths, "scales": scales, - "sid": None, + "sid": torch.tensor([speaker_id]).cpu().numpy(), }, ) return audio[0][0]