diff --git a/tortoise/api.py b/tortoise/api.py index 296ef14a..088c7250 100644 --- a/tortoise/api.py +++ b/tortoise/api.py @@ -1,3 +1,4 @@ +import logging import os import random import uuid @@ -213,7 +214,7 @@ def __init__(self, autoregressive_batch_size=None, models_dir=MODELS_DIR, enable self.models_dir = models_dir self.autoregressive_batch_size = pick_best_batch_size_for_gpu() if autoregressive_batch_size is None else autoregressive_batch_size self.enable_redaction = enable_redaction - self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') + self._set_device(device) if self.enable_redaction: self.aligner = Wav2VecAlignment() @@ -250,6 +251,14 @@ def __init__(self, autoregressive_batch_size=None, models_dir=MODELS_DIR, enable self.rlg_auto = None self.rlg_diffusion = None + def _set_device(self, device): + if device is None: + logging.warning("No device specified. This will default to the first GPU if available, otherwise CPU.") + device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') + if device == 'cpu' or (isinstance(device, torch.device) and device.type == 'cpu'): + logging.warning("Running on CPU. This will be painfully slow. You've been warned.") + self.device = device + def load_cvvp(self): """Load CVVP model.""" self.cvvp = CVVP(model_dim=512, transformer_heads=8, dropout=0, mel_codes=8192, conditioning_enc_depth=8, cond_mask_percentage=0,