import torchaudio as ta
import torch
from chatterbox.mtl_tts import ChatterboxMultilingualTTS
import functools

# torch.load = functools.partial(torch.load, map_location='cpu')

# device_map = torch.device('cpu')
device_map = None
if torch.cuda.is_available():
    device_map = torch.device('cuda')
else:
    device_map = torch.device('cpu')

print(f"Using device: {device_map}")

tts_model = ChatterboxMultilingualTTS.from_pretrained(device=device_map)
streamer_lang = "es"

msg = "CDOM201 dice: Como estas pandita, igual de puto como siempre?"
audio_file = tts_model.generate(msg, language_id=streamer_lang)

ta.save("sleeplespanda.wav", audio_file, tts_model.sr);