From fb9addb8c0eebde09ec6ba8b81f33dfa0f4dc572 Mon Sep 17 00:00:00 2001
From: BuildTools <chenxuzhang45@gmail.com>
Date: Sat, 31 Aug 2024 13:54:24 -0700
Subject: [PATCH] feat(models): add HF upload/download class

---
 requirements.txt  |  4 +++-
 src/HFTransfer.py | 30 ++++++++++++++++++++++++++++++
 2 files changed, 33 insertions(+), 1 deletion(-)
 create mode 100644 src/HFTransfer.py

diff --git a/requirements.txt b/requirements.txt
index d08783b..8e81420 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -8,4 +8,6 @@ pynvml~=11.5.3
 PySide6~=6.7.2
 flask~=3.0.3
 python-dotenv~=1.0.1
-safetensors~=0.4.4
\ No newline at end of file
+safetensors~=0.4.4
+setuptools~=68.2.0
+huggingface-hub~=0.24.6
diff --git a/src/HFTransfer.py b/src/HFTransfer.py
new file mode 100644
index 0000000..9c08611
--- /dev/null
+++ b/src/HFTransfer.py
@@ -0,0 +1,30 @@
+from huggingface_hub import HfApi, snapshot_download, hf_hub_url
+
+
+class HFTransfer:
+    def __init__(self, token) -> None:
+        self.api = HfApi(token=token)
+
+    def upload_file(self, path, remote_path, repo, repo_type) -> str:
+        status = self.api.upload_file(
+            path_or_fileobj=path,
+            path_in_repo=remote_path,
+            repo_id=repo,
+            repo_type=repo_type,
+        )
+        return status
+
+    def upload_folder(self, local_folder, remote_path, repo, repo_type) -> str:
+        status = self.api.upload_folder(
+            folder_path=local_folder,
+            path_in_repo=remote_path,
+            repo_id=repo,
+            repo_type=repo_type,
+        )
+        return status
+
+    def download_repo(self, repo) -> str:
+        return snapshot_download(repo_id=repo)
+
+    def get_download_link(self, repo, file) -> str:
+        return hf_hub_url(repo_id=repo, filename=file)