[RFC][Feature request] Loading model onto specific GPU

[jaketae]

Context

The general API for TTS inference looks like this:

from TTS.api import tts
model = tts("tts_models/en/ljspeech/glow-tts", gpu=True)

gpu is a bool, so users can only specify True or False.

Problem

If the user is trying to load the model on a machine with multiple GPU cards, perhaps they might want to do something like

model = tts("tts_models/en/ljspeech/glow-tts", gpu="cuda:1")

Although running a script with the CUDA_VISIBLE_DEVICES is an option, this could be limiting if the user is dealing with loading multiple models, and they want the specific TTS model to go to cuda:x and another NLP model to go to cuda:y, for example.

Proposed Solution

Tracing the gpu flag, it appears to be used in places like synthesizer.py:

TTS/TTS/utils/synthesizer.py

Lines 118 to 119 in 14d45b5

	if use_cuda:
	self.tts_model.cuda()

Instead of calling .cuda(), perhaps we can make the flag more flexible, say device, and call .to(device) instead. Obviously, changing the name of the argument would be breaking, so to preserve BC, we should introduce this as another optional keyword argument.

I'm not too familiar with the details of the TTS API, so any feedback would be appreciated!

[stale]

This issue has been automatically marked as stale because it has not had recent activity. It will be closed if no further activity occurs. Thank you for your contributions. You might also look our discussion channels.

[jaketae]

Unstale

[stale]

This issue has been automatically marked as stale because it has not had recent activity. It will be closed if no further activity occurs. Thank you for your contributions. You might also look our discussion channels.

[Edresson]

Hi @jaketae,

You can easily select the device using the "CUDA_VISIBLE_DEVICES=1". Using this only GPU 1 will be visible for the Pytorch and it will be forced to use this GPU. Example:

CUDA_VISIBLE_DEVICES=1 python3  api.py 

You can also use the "torch.cuda.set_device(device_num)" to set the device however I recommend you to use "CUDA_VISIBLE_DEVICES".

[jaketae]

@Edresson, thanks for the reply.

I'm aware of CUDA_VISIBLE_DEVICES. However, this sets the visibility of GPUs on the script level. If the user is loading multiple models onto the GPU in a single script (e.g., a language model, TTS model, and an image generation model), they might want to load these models to different GPU cards. Using CUDA_VISIBLE_DEVICES would force the script to load all models to a single GPU, which is not the desired behavior.

For context, I am building StoryTeller, a package that leverages TTS, diffusion, and GPT to create a narrated video of a story. I was wondering if there is a way to specify to which cuda:x device a TTS module can be loaded. One can imagine users loading a language model to cuda:0, TTS to cuda:1, and Stable Diffusion to cuda:2, for instance. This doesn't seem to be possible with the current design.

If this is not supported in CoquiTTS, I'm happy to open a PR. My initial thought is to modify this line

TTS/TTS/utils/synthesizer.py

Lines 118 to 119 in 14d45b5

	if use_cuda:
	self.tts_model.cuda()

but if maintainers have other thoughts or design philosophies, I'm curious to hear them as well.

[erogol]

I just reopen this if anyone is willing to contribute

[jaketae]

@erogol @Edresson, thanks for the reply.

I'm interested in contributing. Some questions:

1. Is there a policy on deprecation or backward compatibility? One way is using a device argument instead of gpu, but getting rid of the gpu argument would break BC. In this case, should we keep the flag and maybe raise a warning that gpu will be deprecated in future releases?
2. Alternatively, we could add a .to(device) method directly to the TTS or Synthesizer class, either by creating these methods or by subclassing from nn.Module. Do you have a preference as to how we go about this feature?

Once there is some agreement around the design, I'll submit a PR soon. Thanks!

[erogol]

thanks @jaketae for you interest

I like the second option. I'd not worry about BC since we have not released a major version by definition we don't promise BC. But of course for convenience it'd be nice to do the transition with a warning.

[stale]

This issue has been automatically marked as stale because it has not had recent activity. It will be closed if no further activity occurs. Thank you for your contributions. You might also look our discussion channels.

[stale]

This issue has been automatically marked as stale because it has not had recent activity. It will be closed if no further activity occurs. Thank you for your contributions. You might also look our discussion channels.

[jaketae]

Apologies for the delay; I was carried away by other work and did not get the chance to work on this.

I can still give it a try in the next few weeks, but if anyone else wants to give it a shot, feel free!

[erogol]

Feel free to send a PR whenever.

[jaketae]

WIP draft PR in #2855.