phoonnx 0.2.3__py3-none-any.whl → 0.2.4__py3-none-any.whl

phoonnx/version.py

@@ -1,7 +1,7 @@
 # START_VERSION_BLOCK
 VERSION_MAJOR = 0
 VERSION_MINOR = 2
-VERSION_BUILD = 3
+VERSION_BUILD = 4
 VERSION_ALPHA = 0
 # END_VERSION_BLOCK
 

{phoonnx-0.2.3.dist-info → phoonnx-0.2.4.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: phoonnx
-Version: 0.2.3
+Version: 0.2.4
 Home-page: https://github.com/TigreGotico/phoonnx
 Author: JarbasAi
 Author-email: jarbasai@mailfence.com

{phoonnx-0.2.3.dist-info → phoonnx-0.2.4.dist-info}/RECORD

@@ -2,7 +2,7 @@ phoonnx/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 phoonnx/config.py,sha256=DKgsU03g8jrAuMcVqbu-w3MWPXOUihFtRnavg6WGQ1Y,19983
 phoonnx/phoneme_ids.py,sha256=FiNgZwV6naEsBh6XwFLh3_FyOgPiCsK9qo7S0v-CmI4,13667
 phoonnx/util.py,sha256=XSjFEoqSFcujFTHxednacgC9GrSYyF-Il5L6Utmxmu4,25909
-phoonnx/version.py,sha256=pQLK9MT9qYgrDxum7k-V4nQPl8az4esj-UN-cAHGOKU,237
+phoonnx/version.py,sha256=28DZfjsylGY3KkroBs-6sispnnXZDH3ZHCzHcKoK5L0,237
 phoonnx/voice.py,sha256=JXjmbrhJd4mmTiLgz4O_Pa5_rKGUC9xzuBfqxYDw3Mg,19420
 phoonnx/locale/ca/phonetic_spellings.txt,sha256=igv3t7jxLSRE5GHsdn57HOpxiWNcEmECPql6m02wbO0,47
 phoonnx/locale/en/phonetic_spellings.txt,sha256=xGQlWOABLzbttpQvopl9CU-NnwEJRqKx8iuylsdUoQA,27
@@ -63,8 +63,8 @@ phoonnx/thirdparty/tashkeel/input_id_map.json,sha256=cnpJqjx-k53AbzKyfC4GxMS771l
 phoonnx/thirdparty/tashkeel/model.onnx,sha256=UsQNQsoJT_n_B6CR0KHq_XuqXPI4jmCpzIm6zY5elV8,4788213
 phoonnx/thirdparty/tashkeel/target_id_map.json,sha256=baNAJL_UwP9U91mLt01aAEBRRNdGr-csFB_O6roh7TA,181
 phoonnx_train/export_onnx.py,sha256=CPfgNEm0hnXPSlgme0R9jr-6jZ5fKFpG5DZJFMkC-h4,12820
-phoonnx_train/preprocess.py,sha256=4FJFi7KL-ZUmrbN2NyhxBNpEjDlPRLSDJo2JoyvpR14,21700
-phoonnx_train/train.py,sha256=16HAb6Yu51xmbP3VM-tWUCsIhct1JHf56aQLG2UEaIc,6024
+phoonnx_train/preprocess.py,sha256=T1YcM89bizARKMaqjDtr3JzidoNvAbm-3n3eClbzMsI,21880
+phoonnx_train/train.py,sha256=nsINvDQ3dYvBne5UWPgLZ0a4qZFdSsOKk8HzZHGTLY4,8757
 phoonnx_train/norm_audio/__init__.py,sha256=Al_YwqMnENXRWp0c79cDZqbdd7pFYARXKxCfBaedr1c,3030
 phoonnx_train/norm_audio/trim.py,sha256=_ZsE3SYhahQSdEdBLeSwyFJGcvEbt-5E_lnWwTT4tcY,1698
 phoonnx_train/norm_audio/vad.py,sha256=DXHfRD0qqFJ52FjPvrL5LlN6keJWuc9Nf6TNhxpwC_4,1600
@@ -83,7 +83,7 @@ phoonnx_train/vits/utils.py,sha256=exiyrtPHbnnGvcHWSbaH9-gR6srH5ZPHlKiqV2IHUrQ,4
 phoonnx_train/vits/wavfile.py,sha256=oQZiTIrdw0oLTbcVwKfGXye1WtKte6qK_52qVwiMvfc,26396
 phoonnx_train/vits/monotonic_align/__init__.py,sha256=5IdAOD1Z7UloMb6d_9NRFsXoNIjEQ3h9mvOSh_AtO3k,636
 phoonnx_train/vits/monotonic_align/setup.py,sha256=0K5iJJ2mKIklx6ncEfCQS34skm5hHPiz9vRlQEvevvY,266
-phoonnx-0.2.3.dist-info/METADATA,sha256=xeznPAJFWZSkqGXdS30OaONue2i4YNd12ervgl8sjCw,8248
-phoonnx-0.2.3.dist-info/WHEEL,sha256=tZoeGjtWxWRfdplE7E3d45VPlLNQnvbKiYnx7gwAy8A,92
-phoonnx-0.2.3.dist-info/top_level.txt,sha256=ZrnHXe-4HqbOSX6fbdY-JiP7YEu2Bok9T0ji351MrmM,22
-phoonnx-0.2.3.dist-info/RECORD,,
+phoonnx-0.2.4.dist-info/METADATA,sha256=UcIFJkCXmN-YrP-_QXTIS63sJne3AnVhr_qrrZxmr68,8248
+phoonnx-0.2.4.dist-info/WHEEL,sha256=tZoeGjtWxWRfdplE7E3d45VPlLNQnvbKiYnx7gwAy8A,92
+phoonnx-0.2.4.dist-info/top_level.txt,sha256=ZrnHXe-4HqbOSX6fbdY-JiP7YEu2Bok9T0ji351MrmM,22
+phoonnx-0.2.4.dist-info/RECORD,,

phoonnx_train/preprocess.py

@@ -32,7 +32,7 @@ DEFAULT_SPECIAL_PHONEME_ID_MAP: Dict[str, int] = {
     DEFAULT_EOS_TOKEN: 2,
     DEFAULT_BLANK_WORD_TOKEN: 3,
 }
-
+MAX_PHONEMES = 256
 # -----------------------------------------------------------------------------
 
 @dataclass
@@ -517,12 +517,15 @@ def cli(
 
     if prev_config:
         with open(prev_config) as f:
-            prev_phoneme_id_map = json.load(f)["phoneme_id_map"]
+            cfg = json.load(f)
+        prev_phoneme_id_map = cfg["phoneme_id_map"]
+        prev_num_symbols = cfg.get("num_symbols", MAX_PHONEMES)
         _LOGGER.info(f"Loaded phoneme map from previous config: '{prev_config}'")
         all_phonemes.update(prev_phoneme_id_map.keys())
         final_phoneme_id_map = prev_phoneme_id_map
-        _LOGGER.info("previous phoneme map contains %d symbols.", len(final_phoneme_id_map))
+        _LOGGER.info("previous phoneme map contains %d phonemes.", len(final_phoneme_id_map))
     else:
+        prev_num_symbols = MAX_PHONEMES
         final_phoneme_id_map: Dict[str, int] = DEFAULT_SPECIAL_PHONEME_ID_MAP.copy()
         if phonemizer.alphabet == Alphabet.IPA:
             all_phonemes.update(DEFAULT_IPA_PHONEME_ID_MAP.keys())
@@ -533,7 +536,7 @@ def cli(
                                       if p not in existing_keys]
                                      )
 
-    _LOGGER.info("Collected %d new symbols.", len(new_phonemes))
+    _LOGGER.info("Collected %d new phonemes.", len(new_phonemes))
 
     finetune_error = prev_config and len(new_phonemes)
     if finetune_error:
@@ -553,7 +556,7 @@ def cli(
             _LOGGER.debug(f"New phoneme: {pho}")
 
     if new_phonemes:
-        _LOGGER.info("Final phoneme map contains %d symbols.", len(final_phoneme_id_map))
+        _LOGGER.info("Final phoneme map contains %d phonemes.", len(final_phoneme_id_map))
 
     # --- Write the final config.json ---
     _LOGGER.info("Writing dataset config...")
@@ -575,7 +578,7 @@ def cli(
         "phoneme_type": config.phoneme_type.value,
         "phonemizer_model": config.phonemizer_model,
         "phoneme_id_map": final_phoneme_id_map,
-        "num_symbols": len(final_phoneme_id_map),
+        "num_symbols": prev_num_symbols if prev_config else len(final_phoneme_id_map),
         "num_speakers": len(speaker_counts) if is_multispeaker else 1,
         "speaker_id_map": speaker_ids,
         "phoonnx_version": VERSION_STR,

phoonnx_train/train.py

@@ -44,6 +44,7 @@ def load_state_dict(model, saved_state_dict):
 @click.option('--batch-size', type=int, default=16, help='Training batch size (default: 16)')
 @click.option('--num-workers', type=click.IntRange(min=1), default=os.cpu_count() or 1, help='Number of data loader workers (default: CPU count)')
 @click.option('--validation-split', type=float, default=0.05, help='Proportion of data used for validation (default: 0.05)')
+@click.option('--discard-encoder', type=bool, default=False, help='Discard the encoder weights from base checkpoint (default: False)')
 def main(
     dataset_dir,
     checkpoint_epochs,
@@ -60,6 +61,7 @@ def main(
     batch_size,
     num_workers,
     validation_split,
+    discard_encoder
 ):
     logging.basicConfig(level=logging.DEBUG)
 
@@ -73,22 +75,18 @@ def main(
     config_path = dataset_dir / 'config.json'
     dataset_path = dataset_dir / 'dataset.jsonl'
 
-    print(f"INFO - config_path: '{config_path}'")
-    print(f"INFO - dataset_path: '{dataset_path}'")
+    _LOGGER.info(f"config_path: '{config_path}'")
+    _LOGGER.info(f"dataset_path: '{dataset_path}'")
 
     with open(config_path, 'r', encoding='utf-8') as config_file:
         config = json.load(config_file)
-        num_symbols = int(config['num_symbols'])
-        num_speakers = int(config['num_speakers'])
-        sample_rate = int(config['audio']['sample_rate'])
 
     trainer = Trainer(
         max_epochs=max_epochs,
         devices=devices,
         accelerator=accelerator,
         default_root_dir=default_root_dir,
-        precision=precision,
-        resume_from_checkpoint=resume_from_checkpoint
+        precision=precision
     )
 
     if checkpoint_epochs is not None:
@@ -119,7 +117,32 @@ def main(
             'upsample_kernel_sizes': (16, 16, 4, 4),
         })
 
-    print(f"VitsModel params: num_symbols={num_symbols} num_speakers={num_speakers} sample_rate={sample_rate}")
+    num_symbols = int(config['num_symbols'])
+    num_speakers = int(config['num_speakers'])
+    sample_rate = int(config['audio']['sample_rate'])
+    _LOGGER.debug(f"Config params: num_symbols={num_symbols} num_speakers={num_speakers} sample_rate={sample_rate}")
+
+    if resume_from_checkpoint:
+        # TODO (?) - add a flag to use params from config vs from checkpoint in case of mismatch
+        ckpt = VitsModel.load_from_checkpoint(resume_from_checkpoint, dataset=None)
+        _LOGGER.debug(f"Checkpoint params: num_symbols={ckpt.model_g.n_vocab} num_speakers={ckpt.model_g.n_speakers} sample_rate={ckpt.hparams.sample_rate}")
+        if ckpt.model_g.n_vocab != num_symbols:
+            _LOGGER.warning(f"Checkpoint num_symbols={ckpt.model_g.n_vocab} does not match config num_symbols={num_symbols}")
+            #-------------
+            # commented out this code because this is not supposed to happen if you used the preprocess.py script
+            # uncomment if you want to use the encoder from checkpoint + update num_symbols in the .json file manually
+            #-------------
+            #if ckpt.model_g.n_vocab > num_symbols and not discard_encoder:
+            #    num_symbols = ckpt.model_g.n_vocab
+            #    _LOGGER.info(f"Training with num_symbols={num_symbols}")
+            ###############
+        if ckpt.model_g.n_speakers != num_speakers:
+            _LOGGER.warning(f"Checkpoint num_speakers={ckpt.model_g.n_speakers} does not match config num_speakers={num_speakers}")
+            #num_speakers = ckpt.model_g.n_speakers
+        if ckpt.hparams.sample_rate != sample_rate:
+            _LOGGER.warning(f"Checkpoint sample_rate={ckpt.hparams.sample_rate} does not match config sample_rate={sample_rate}")
+            #sample_rate = ckpt.hparams.sample_rate
+
     model = VitsModel(
         num_symbols=num_symbols,
         num_speakers=num_speakers,
@@ -127,6 +150,31 @@ def main(
         dataset=[dataset_path],
         **dict_args,
     )
+    _LOGGER.info(f"VitsModel params: num_symbols={num_symbols} num_speakers={num_speakers} sample_rate={sample_rate}")
+
+    if resume_from_checkpoint:
+        saved_state_dict = ckpt.state_dict()
+
+        # Filter the state dictionary by removing the encoder weights
+        enc_key = 'model_g.enc_p.emb.weight'
+        if enc_key in saved_state_dict:
+            saved_shape = saved_state_dict[enc_key].shape
+            current_shape = model.state_dict()[enc_key].shape
+            if saved_shape[0] != current_shape[0]:
+                _LOGGER.warning(
+                    "Size mismatch detected for '%s': saved shape %s vs current shape %s. ",
+                    enc_key, saved_shape, current_shape
+                )
+                discard_encoder = True
+
+            if discard_encoder:
+                _LOGGER.warning(
+                    "Skipping encoder weights from the checkpoint. (will be randomly initialized)"
+                )
+                saved_state_dict.pop(enc_key)
+
+        load_state_dict(model, saved_state_dict)
+        _LOGGER.info("Successfully loaded model weights.")
 
     if resume_from_single_speaker_checkpoint:
         assert num_speakers > 1, "--resume-from-single-speaker-checkpoint is only for multi-speaker models."
@@ -143,7 +191,7 @@ def main(
         load_state_dict(model.model_d, model_single.model_d.state_dict())
         _LOGGER.info('Successfully converted single-speaker checkpoint to multi-speaker')
 
-    print('training started!!')
+    _LOGGER.info('training started!!')
     trainer.fit(model)