|
from pathlib import Path
|
|
from typing import Any, Optional, Union
|
|
|
|
import torch
|
|
from safetensors import safe_open
|
|
from safetensors.torch import save_file
|
|
|
|
from style_bert_vits2.logging import logger
|
|
|
|
|
|
def load_safetensors(
|
|
checkpoint_path: Union[str, Path],
|
|
model: torch.nn.Module,
|
|
for_infer: bool = False,
|
|
) -> tuple[torch.nn.Module, Optional[int]]:
|
|
"""
|
|
指定されたパスから safetensors モデルを読み込み、モデルとイテレーションを返す。
|
|
|
|
Args:
|
|
checkpoint_path (Union[str, Path]): モデルのチェックポイントファイルのパス
|
|
model (torch.nn.Module): 読み込む対象のモデル
|
|
for_infer (bool): 推論用に読み込むかどうかのフラグ
|
|
|
|
Returns:
|
|
tuple[torch.nn.Module, Optional[int]]: 読み込まれたモデルとイテレーション回数（存在する場合）
|
|
"""
|
|
|
|
tensors: dict[str, Any] = {}
|
|
iteration: Optional[int] = None
|
|
with safe_open(str(checkpoint_path), framework="pt", device="cpu") as f:
|
|
for key in f.keys():
|
|
if key == "iteration":
|
|
iteration = f.get_tensor(key).item()
|
|
tensors[key] = f.get_tensor(key)
|
|
if hasattr(model, "module"):
|
|
result = model.module.load_state_dict(tensors, strict=False)
|
|
else:
|
|
result = model.load_state_dict(tensors, strict=False)
|
|
for key in result.missing_keys:
|
|
if key.startswith("enc_q") and for_infer:
|
|
continue
|
|
logger.warning(f"Missing key: {key}")
|
|
for key in result.unexpected_keys:
|
|
if key == "iteration":
|
|
continue
|
|
logger.warning(f"Unexpected key: {key}")
|
|
if iteration is None:
|
|
logger.info(f"Loaded '{checkpoint_path}'")
|
|
else:
|
|
logger.info(f"Loaded '{checkpoint_path}' (iteration {iteration})")
|
|
|
|
return model, iteration
|
|
|
|
|
|
def save_safetensors(
|
|
model: torch.nn.Module,
|
|
iteration: int,
|
|
checkpoint_path: Union[str, Path],
|
|
is_half: bool = False,
|
|
for_infer: bool = False,
|
|
) -> None:
|
|
"""
|
|
モデルを safetensors 形式で保存する。
|
|
|
|
Args:
|
|
model (torch.nn.Module): 保存するモデル
|
|
iteration (int): イテレーション回数
|
|
checkpoint_path (Union[str, Path]): 保存先のパス
|
|
is_half (bool): モデルを半精度で保存するかどうかのフラグ
|
|
for_infer (bool): 推論用に保存するかどうかのフラグ
|
|
"""
|
|
|
|
if hasattr(model, "module"):
|
|
state_dict = model.module.state_dict()
|
|
else:
|
|
state_dict = model.state_dict()
|
|
keys = []
|
|
for k in state_dict:
|
|
if "enc_q" in k and for_infer:
|
|
continue
|
|
keys.append(k)
|
|
|
|
new_dict = (
|
|
{k: state_dict[k].half() for k in keys}
|
|
if is_half
|
|
else {k: state_dict[k] for k in keys}
|
|
)
|
|
new_dict["iteration"] = torch.LongTensor([iteration])
|
|
logger.info(f"Saved safetensors to {checkpoint_path}")
|
|
|
|
save_file(new_dict, checkpoint_path)
|
|
|
