Update hf_merge.py

hf_merge.py

@@ -121,19 +121,33 @@ class ModelMerger:
             return True # for safe
     
     
-    def upload_model(new_repo_id, diffusers_folder, is_private, token):
-        from huggingface_hub import HfApi
-        api = HfApi(token=token)
-        try:
-            api.create_repo(repo_id=new_repo_id, token=token)
-            api.upload_folder(repo_id=new_repo_id, folder_path=output_dir, token=token)
-            url = f"https://huggingface.co/{repo_name}"
-        except Exception as e:
-            print(f"Error: Failed to upload to {repo_name}.")
-            print(e)
-            return ""
-        return url
-
+    def upload_model(self, output_dir, commit_message):
+            """
+            Upload the merged model to HuggingFace.
+    
+            Args:
+                output_dir (str): The path to the output directory containing the merged model.
+                commit_message (str): The commit message for the upload.
+            """
+            try:
+                # Create a new repository if it doesn't exist
+                if not self.api.repo_exists(repo_id=self.repo_id, token=self.token):
+                    self.api.create_repo(repo_id=self.repo_id, token=self.token, private=True)
+    
+                # Upload the folder to the repository
+                self.api.upload_folder(
+                    repo_id=self.repo_id,
+                    folder_path=output_dir,
+                    commit_message=commit_message,
+                    token=self.token
+                )
+                url = f"https://huggingface.co/{self.repo_id}"
+                logging.info(f"Model uploaded successfully to {url}")
+                return url
+            except Exception as e:
+                logging.error(f"Error: Failed to upload to {self.repo_id}.")
+                logging.error(e)
+                return ""
 def get_max_vocab_size(repo_list):
     """
     Get the maximum vocabulary size from a list of repositories.
@@ -211,7 +225,13 @@ def main():
         if not args.token or not args.repo:
             logging.error("Error: HuggingFace token and repo name are required for uploading.")
         else:
-            model_merger.upload_model(output_dir, args.repo, args.commit_message)
+            url = model_merger.upload_model(output_dir, args.commit_message)
+            if url:
+                logging.info(f"Model uploaded successfully to {url}")
+
+if __name__ == "__main__":
+    main()
+
 
 if __name__ == "__main__":
     main()