Fix hf push

train_ms.py

@@ -6,6 +6,7 @@ import platform
 
 import torch
 import torch.distributed as dist
+from huggingface_hub import HfApi
 from torch.cuda.amp import GradScaler, autocast
 from torch.nn import functional as F
 from torch.nn.parallel import DistributedDataParallel as DDP
@@ -45,6 +46,8 @@ torch.backends.cuda.enable_math_sdp(True)
 
 global_step = 0
 
+api = HfApi()
+
 
 def run():
     # Command line configuration is not recommended unless necessary, use config.yml
@@ -483,6 +486,18 @@ def run():
                 ),
                 for_infer=True,
             )
+            if hps.repo_id is not None:
+                api.upload_folder(
+                    repo_id=hps.repo_id,
+                    folder_path=model_dir,
+                    path_in_repo=f"Data/{config.model_name}/models",
+                    delete_patterns="*.pth",
+                )
+                api.upload_folder(
+                    repo_id=hps.repo_id,
+                    folder_path=config.out_dir,
+                    path_in_repo=f"model_assets/{config.model_name}",
+                )
 
     if pbar is not None:
         pbar.close()
@@ -765,6 +780,18 @@ def train_and_evaluate(
                     ),
                     for_infer=True,
                 )
+                if hps.repo_id is not None:
+                    api.upload_folder(
+                        repo_id=hps.repo_id,
+                        folder_path=hps.model_dir,
+                        path_in_repo=f"Data/{config.model_name}/models",
+                        delete_patterns="*.pth",
+                    )
+                    api.upload_folder(
+                        repo_id=hps.repo_id,
+                        folder_path=config.out_dir,
+                        path_in_repo=f"model_assets/{config.model_name}",
+                    )
 
         global_step += 1
         if pbar is not None:

train_ms_jp_extra.py

@@ -578,10 +578,13 @@ def run():
                 api.upload_folder(
                     repo_id=hps.repo_id,
                     folder_path=model_dir,
+                    path_in_repo=f"Data/{config.model_name}/models",
+                    delete_patterns="*.pth",
                 )
                 api.upload_folder(
                     repo_id=hps.repo_id,
                     folder_path=config.out_dir,
+                    path_in_repo=f"model_assets/{config.model_name}",
                 )
 
     if pbar is not None:
@@ -937,11 +940,14 @@ def train_and_evaluate(
                 if hps.repo_id is not None:
                     api.upload_folder(
                         repo_id=hps.repo_id,
-                        folder_path=model_dir,
+                        folder_path=hps.model_dir,
+                        path_in_repo=f"Data/{config.model_name}/models",
+                        delete_patterns="*.pth",
                     )
                     api.upload_folder(
                         repo_id=hps.repo_id,
                         folder_path=config.out_dir,
+                        path_in_repo=f"model_assets/{config.model_name}",
                     )
 
         global_step += 1