import torch
from TTS.tts.models.xtts import XttsAudioConfig  
from TTS.tts.configs.xtts_config import XttsConfig  
from TTS.config.shared_configs import BaseDatasetConfig  
from TTS.tts.models.xtts import XttsArgs  

# **添加 `XttsConfig` 到 PyTorch 安全全局对象**
torch.serialization.add_safe_globals([XttsConfig, XttsAudioConfig, BaseDatasetConfig, XttsArgs])

# **强制 `weights_only=False` 解决 UnpicklingError**
# def load_fsspec_fixed(*args, **kwargs):
#     kwargs["weights_only"] = False  # 关键修正
#     return torch.load(*args, **kwargs)

def load_fsspec_fixed(*args, **kwargs):
    kwargs.pop("cache", None)  # 移除 cache 参数，避免报错
    kwargs["weights_only"] = False  # 关键修正
    return torch.load(*args, **kwargs)


# **覆盖 Coqui TTS 的 `load_fsspec` 方法**
import TTS.utils.io
TTS.utils.io.load_fsspec = load_fsspec_fixed

from TTS.api import TTS

# **XTTS v2 高质量多语言模型**
XTTS_MODEL = "tts_models/multilingual/multi-dataset/xtts_v2"

# **VITS 高质量单语言模型**
VITS_MODEL = "tts_models/en/ljspeech/vits"

# **高质量的参考音频**
speaker_wav_en = "example_speaker_en.wav"  # 必须提供此音频，确保质量良好

# **要转换的文本**
text_en = "This is a high-quality text-to-speech conversion using XTTS v2 and VITS."
text = "记者从越秀区了解到，广州博物馆隆重推出“吉祥有年——广州博物馆藏吉祥文物展”及“吉祥有年·潮派趁墟”主题新春市集，让市民群众在探秘文物珍宝的同时，也能解锁非遗新体验、品尝地道广府味，一起喜迎吉祥乙巳蛇年。"

# **加载 XTTS v2**
print("🚀 正在加载 XTTS v2 模型，请稍候...")
tts_xtts = TTS(XTTS_MODEL)

# **使用 XTTS v2 生成语音**
tts_xtts.tts_to_file(
    text=text, 
    file_path="output_xtts.wav", 
    speaker_wav=speaker_wav_en,  # 参考音频（用于克隆音色）
    language="zh-cn",  # 语言代码，必须匹配文本语言
    split_sentences=True  # 让模型自动优化长文本
)
print("✅ XTTS v2 语音合成完成！已保存到 output_xtts.wav 🎵")

# **加载 VITS**
print("🚀 正在加载 VITS 模型，请稍候...")
tts_vits = TTS(VITS_MODEL)

# **使用 VITS 生成语音**
tts_vits.tts_to_file(
    text=text_en,
    split_sentences=True,
    file_path="output_vits.wav"
)
print("✅ VITS 语音合成完成！已保存到 output_vits.wav 🎵")