Add device support in TTS and Synthesizer

[jaketae]

Context

In #2282, we proposed up the possibility of implementing atts.to(device) interface as a substitute for use_cuda or gpu flags. The current flags do not allow users to specify the specific GPU device (e.g., cuda:3). It also does not allow users to use other accelerated backends, such as Apple Silicon GPUs (MPS), which PyTorch now supports.

Solution

We make TTS and Synthesizer classes inherit from nn.Module. This gives us .to(device) for free for both of the classes.

We can now run TTS on Apple Silicon (tested on M2 Max). Not all kernels have been implemented in MPS in PyTorch yet, so we need to set the environment variable

export PYTORCH_ENABLE_MPS_FALLBACK=1

to enable CPU fallback. With this set, we can now run

>>> from TTS.api import TTS
>>> model_name = TTS.list_models()[0]
>>> tts = TTS(model_name)
>>> tts = tts.to("mps")
>>> tts.tts_to_file(text="Hello world!", speaker=tts.speakers[0], language=tts.languages[0], file_path="output.wav")

Also tested with make test.

[CLAassistant]

All committers have signed the CLA.

[jaketae]

Added warning. We could add specific dates or versions to better inform users about future plans, but I left it this way because I didn't have enough context on the future releases roadmap.

[jaketae]

Added new device argument to functions called in Synthesizer. To retain backwards compatibility, we keep the cuda argument for now; we should probably clean them up in the future and provide a single way of configuring device/enabling CUDA.

[jaketae]

In some obscure use cases, the user could have placed the feature frontend and the vocoder on different devices.

>>> tts.synthesizer.tts_model = tts.synthesizer.tts_model.to("cuda:0")
>>> tts.synthesizer.vocoder_model = tts.synthesizer.vocoder_model.to("cuda:1")

We check the device of the vocoder, if it exists.

[jaketae]

Hi @erogol, curious to hear your thoughts on this implementation! The guiding philosophy was to use the PyTorch .to(device) API while keeping all functionality intact and retaining backwards compatibility with use_cuda or gpu.

---

I've signed the CLA, but the first few commits didn't have my GitHub email (I just got a new laptop and forgot to set up my Git user information), which is why the CLA test is marked as pending.

[erogol]

@jaketae thanks for the PR. I'll review it tomorrow 👍

[erogol]

All looks good!! Thanks for the PR. If you think its done I can merge.

[jaketae]

@erogol, thanks for the quick review! I think we can go ahead with the merge unless you have second thoughts. I'll maybe open a follow-up PR to improve docs or the README where applicable. Thanks!

[erogol]

@jaketae awesome thanks again