From e0a601737a1a457810ea28618993f503a0eb57bf Mon Sep 17 00:00:00 2001
From: litagin02 <litagin02@gmail.com>
Date: Sat, 17 Feb 2024 20:26:58 +0900
Subject: [PATCH] Error handling for hug upload

---
 train_ms.py          | 50 +++++++++++++++++++++++++-------------------
 train_ms_jp_extra.py | 50 +++++++++++++++++++++++++-------------------
 2 files changed, 56 insertions(+), 44 deletions(-)

diff --git a/train_ms.py b/train_ms.py
index 08d6af9..9fbc991 100644
--- a/train_ms.py
+++ b/train_ms.py
@@ -487,17 +487,20 @@ 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}",
-                )
+                try:
+                    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}",
+                    )
+                except Exception as e:
+                    logger.warning(e)
 
     if pbar is not None:
         pbar.close()
@@ -781,17 +784,20 @@ 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}",
-                    )
+                    try:
+                        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}",
+                        )
+                    except Exception as e:
+                        logger.warning(e)
 
         global_step += 1
         if pbar is not None:
diff --git a/train_ms_jp_extra.py b/train_ms_jp_extra.py
index 9c116ab..ac16661 100644
--- a/train_ms_jp_extra.py
+++ b/train_ms_jp_extra.py
@@ -575,17 +575,20 @@ 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}",
-                )
+                try:
+                    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}",
+                    )
+                except Exception as e:
+                    logger.warning(e)
 
     if pbar is not None:
         pbar.close()
@@ -938,17 +941,20 @@ 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}",
-                    )
+                    try:
+                        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}",
+                        )
+                    except Exception as e:
+                        logger.warning(e)
 
         global_step += 1
         if pbar is not None: