ADJUST io.py for logging mlflow; Credits: Blirona

trainer/io.py

@@ -6,7 +6,7 @@
 from pathlib import Path
 from typing import Any, Callable, Dict, List, Tuple, Union
 from urllib.parse import urlparse
-
+import mlflow
 import fsspec
 import torch
 from coqpit import Coqpit
@@ -161,6 +161,15 @@ def save_checkpoint(
         save_func=save_func,
         **kwargs,
     )
+    from TTS.tts.models.modelWrapper import MyModel
+    from TTS.utils.synthesizer import Synthesizer
+    synthesizer = Synthesizer(checkpoint_path, os.path.join(output_folder, 'config.json'))
+    model = MyModel(synthesizer=synthesizer)
+    # For MLflow logging, , artifacts={'checkpoint': file_name, "output": output_folder}
+
+    run_name = os.getenv("RUN_NAME", "vits_janika_de")
+
+    mlflow.pyfunc.log_model(python_model=model, artifact_path="models/TTS",registered_model_name=run_name, code_path=[output_folder])
     if save_n_checkpoints is not None:
         keep_n_checkpoints(output_folder, save_n_checkpoints)