diff --git a/bert_gen.py b/bert_gen.py index 49be3e6..279b2bb 100644 --- a/bert_gen.py +++ b/bert_gen.py @@ -55,7 +55,7 @@ def process_line(x: tuple[str, bool]): bert_path = wav_path.replace(".WAV", ".wav").replace(".wav", ".bert.pt") try: - bert = torch.load(bert_path) + bert = torch.load(bert_path, weights_only=True) assert bert.shape[-1] == len(phone) except Exception: bert = extract_bert_feature(text, word2ph, Languages(language_str), device) diff --git a/data_utils.py b/data_utils.py index 54d4e27..1bb6364 100644 --- a/data_utils.py +++ b/data_utils.py @@ -129,7 +129,7 @@ class TextAudioSpeakerLoader(torch.utils.data.Dataset): if self.use_mel_spec_posterior: spec_filename = spec_filename.replace(".spec.pt", ".mel.pt") try: - spec = torch.load(spec_filename) + spec = torch.load(spec_filename, weights_only=True) except: if self.use_mel_spec_posterior: spec = mel_spectrogram_torch( @@ -168,7 +168,7 @@ class TextAudioSpeakerLoader(torch.utils.data.Dataset): word2ph[0] += 1 bert_path = wav_path.replace(".wav", ".bert.pt") try: - bert_ori = torch.load(bert_path) + bert_ori = torch.load(bert_path, weights_only=True) assert bert_ori.shape[-1] == len(phone) except Exception as e: logger.warning("Bert load Failed") diff --git a/style_bert_vits2/models/utils/checkpoints.py b/style_bert_vits2/models/utils/checkpoints.py index 756586c..26a8cfa 100644 --- a/style_bert_vits2/models/utils/checkpoints.py +++ b/style_bert_vits2/models/utils/checkpoints.py @@ -32,7 +32,9 @@ def load_checkpoint( """ assert os.path.isfile(checkpoint_path) - checkpoint_dict = torch.load(checkpoint_path, map_location=device) + checkpoint_dict = torch.load( + checkpoint_path, map_location=device, weights_only=True + ) iteration = checkpoint_dict["iteration"] learning_rate = checkpoint_dict["learning_rate"] logger.info( diff --git a/style_bert_vits2/nlp/japanese/user_dict/word_model.py b/style_bert_vits2/nlp/japanese/user_dict/word_model.py index 43420da..8a6780d 100644 --- a/style_bert_vits2/nlp/japanese/user_dict/word_model.py +++ b/style_bert_vits2/nlp/japanese/user_dict/word_model.py @@ -9,7 +9,13 @@ from enum import Enum from re import findall, fullmatch from typing import List, Optional -from pydantic import BaseModel, Field, validator +from pydantic import ( + BaseModel, + ConfigDict, + Field, + ValidationInfo, + field_validator, +) USER_DICT_MIN_PRIORITY = 0 @@ -39,11 +45,11 @@ class UserDictWord(BaseModel): mora_count: Optional[int] = Field(title="モーラ数", default=None) accent_associative_rule: str = Field(title="アクセント結合規則") - class Config: - validate_assignment = True + model_config = ConfigDict(validate_assignment=True, validate_default=True) - @validator("surface") - def convert_to_zenkaku(cls, surface): + @field_validator("surface") + @classmethod + def convert_to_zenkaku(cls, surface: str) -> str: return surface.translate( str.maketrans( "".join(chr(0x21 + i) for i in range(94)), @@ -51,8 +57,9 @@ class UserDictWord(BaseModel): ) ) - @validator("pronunciation", pre=True) - def check_is_katakana(cls, pronunciation): + @field_validator("pronunciation", mode="before") + @classmethod + def check_is_katakana(cls, pronunciation: str) -> str: if not fullmatch(r"[ァ-ヴー]+", pronunciation): raise ValueError("発音は有効なカタカナでなくてはいけません。") sutegana = ["ァ", "ィ", "ゥ", "ェ", "ォ", "ャ", "ュ", "ョ", "ヮ", "ッ"] @@ -75,9 +82,12 @@ class UserDictWord(BaseModel): ) return pronunciation - @validator("mora_count", pre=True, always=True) - def check_mora_count_and_accent_type(cls, mora_count, values): - if "pronunciation" not in values or "accent_type" not in values: + @field_validator("mora_count", mode="before") + @classmethod + def check_mora_count_and_accent_type( + cls, mora_count: Optional[int], info: ValidationInfo + ) -> Optional[int]: + if "pronunciation" not in info.data or "accent_type" not in info.data: # 適切な場所でエラーを出すようにする return mora_count @@ -91,14 +101,14 @@ class UserDictWord(BaseModel): mora_count = len( findall( f"(?:{rule_others}|{rule_line_i}|{rule_line_u}|{rule_one_mora})", - values["pronunciation"], + info.data["pronunciation"], ) ) - if not 0 <= values["accent_type"] <= mora_count: + if not 0 <= info.data["accent_type"] <= mora_count: raise ValueError( "誤ったアクセント型です({})。 expect: 0 <= accent_type <= {}".format( - values["accent_type"], mora_count + info.data["accent_type"], mora_count ) ) return mora_count diff --git a/tests/test_library_compatibility.py b/tests/test_library_compatibility.py new file mode 100644 index 0000000..578fb3e --- /dev/null +++ b/tests/test_library_compatibility.py @@ -0,0 +1,58 @@ +from pathlib import Path + +import pytest +import torch +from pydantic import ValidationError + +from style_bert_vits2.models.utils.checkpoints import load_checkpoint, save_checkpoint +from style_bert_vits2.nlp.japanese.user_dict.word_model import UserDictWord + + +def _user_dict_word(**overrides: object) -> UserDictWord: + values = { + "surface": "test", + "priority": 5, + "part_of_speech": "名詞", + "part_of_speech_detail_1": "一般", + "part_of_speech_detail_2": "*", + "part_of_speech_detail_3": "*", + "inflectional_type": "*", + "inflectional_form": "*", + "stem": "*", + "yomi": "テスト", + "pronunciation": "テスト", + "accent_type": 1, + "accent_associative_rule": "*", + } + values.update(overrides) + return UserDictWord.model_validate(values) + + +def test_user_dict_word_uses_pydantic_v2_validation() -> None: + word = _user_dict_word() + + assert word.surface == "ｔｅｓｔ" + assert word.mora_count == 3 + + with pytest.raises(ValidationError): + _user_dict_word(pronunciation="てすと") + + +def test_checkpoint_round_trip_with_safe_torch_load(tmp_path: Path) -> None: + checkpoint_path = tmp_path / "checkpoint.pth" + source = torch.nn.Linear(2, 1) + optimizer = torch.optim.AdamW(source.parameters()) + save_checkpoint(source, optimizer, 1e-4, 12, checkpoint_path) + + destination = torch.nn.Linear(2, 1) + destination_optimizer = torch.optim.AdamW(destination.parameters()) + _, _, learning_rate, iteration = load_checkpoint( + checkpoint_path, destination, destination_optimizer + ) + + assert learning_rate == 1e-4 + assert iteration == 12 + for source_parameter, destination_parameter in zip( + source.parameters(), destination.parameters() + ): + assert torch.equal(source_parameter, destination_parameter)