from src.chatterbox.mtl_tts import ChatterboxMultilingualTTS
import torchaudio as ta
from safetensors.torch import load_file as load_safetensors

device = "cpu" # or mps or cuda

multilingual_model = ChatterboxMultilingualTTS.from_pretrained(device=device)

# ----
# Then download the file from huggingface and place it in the current directory.
# ----
t3_state = load_safetensors("Thomcles/Chatterbox-TTS-Persian-Farsi/t3_fa.safetensors", device="cpu")
multilingual_model.t3.load_state_dict(t3_state)
multilingual_model.t3.to(device).eval()

persian_text = "سلام! به آزمایش تبدیل متن به گفتار خوش آمدید."
wav_persian = multilingual_model.generate(persian_text, language_id=None)
ta.save("test-fa.wav", wav_persian, multilingual_model.sr)