VoiceDialogue / src /voice_dialogue /tts /weights_migration.py
hzeng412
Make Qwen3-ASR the default on main; bump to 1.2.0
3c70498
Raw
History Blame Contribute Delete
1.75 kB
"""TTS 预训练权重 safetensors 迁移。
transformers >= 4.56 的安全策略 (CVE-2025-32434) 拒绝在 torch < 2.6 上加载
pytorch_model.bin。transformers 加载时优先使用 model.safetensors,因此首次
启动时把 .bin 转换一次即可,无需升级 torch。
"""
from pathlib import Path
from voice_dialogue.config import paths
from voice_dialogue.utils.logger import logger
PRETRAINED_DIRS = [
"chinese-roberta-wwm-ext-large",
"chinese-hubert-base",
]
def ensure_safetensors_weights() -> None:
"""确保 MoYoYo TTS 的预训练权重存在 safetensors 版本,缺失时从 .bin 转换。"""
moyoyo_path = Path(paths.TTS_MODELS_PATH) / "moyoyo"
for dirname in PRETRAINED_DIRS:
model_dir = moyoyo_path / dirname
bin_path = model_dir / "pytorch_model.bin"
st_path = model_dir / "model.safetensors"
if st_path.exists() or not bin_path.exists():
continue
logger.info(f"[INFO] 首次启动:转换 {dirname} 权重为 safetensors...")
try:
import torch
from safetensors.torch import save_file
state_dict = torch.load(bin_path, map_location="cpu", weights_only=True)
# clone 断开共享内存,safetensors 不允许张量间共享存储
state_dict = {
key: value.clone().contiguous()
for key, value in state_dict.items()
if hasattr(value, "clone")
}
save_file(state_dict, st_path, metadata={"format": "pt"})
logger.info(f"[INFO] {dirname} 转换完成: {st_path.stat().st_size // 1024 ** 2} MB")
except Exception as e:
logger.error(f"[ERROR] 转换 {dirname} 权重失败: {e}")