29 - Added possibility to check for access for a certain model (pre-d…

…ownload step) Signed-off-by: marijnvg-tng <[email protected]> Signed-off-by: marijnvg-tng <[email protected]>
TNG · Nov 26, 2024 · 3d11d37 · 3d11d37
1 parent c5bd31f
commit 3d11d37
Show file tree

Hide file tree

Showing 3 changed files with 41 additions and 0 deletions.
diff --git a/WebUI/src/components/DownloadDialog.vue b/WebUI/src/components/DownloadDialog.vue
@@ -209,6 +209,7 @@ async function showConfirm(downList: DownloadModelParam[], success?: () => void,
     } catch (ex) {
         fail && fail({ type: "error", error: ex });
     }
+  check_model_access()
 }
 
 function getInfoUrl(repoId: string, type: number) {
@@ -280,6 +281,22 @@ function download() {
     })
 }
 
+async function check_model_access() {
+  const response = await fetch(`${globalSetup.apiHost}/api/checkModelAccess`, {
+    method: "POST",
+    body: JSON.stringify([downloadList.value[0].repo_id, models.hfToken]),
+    headers: {
+      "Content-Type": "application/json"
+    }
+  })
+  const data = await response.json()
+  console.log("Is URL valid:")
+  console.log(data.valid)
+  console.log(data.url)
+  console.log(data.status)
+  return data.valid
+}
+
 function cancelConfirm() {
     downloadReject && downloadReject({ type: "cancelConfrim" });
     emits("close");

diff --git a/service/model_downloader.py b/service/model_downloader.py
@@ -301,6 +301,18 @@ def init_download(self, file: HFDonloadItem):
 
         return response, fw
 
+    def is_token_valid(self, repo_id: str):
+        headers={}
+        if (self.hf_token is not None):
+            headers["Authorization"] = f"Bearer {self.hf_token}"
+
+        name = self.fs.ls(repo_id, detail=True)[0].get("name")
+        url = hf_hub_url(repo_id=repo_id, filename = path.basename(path.relpath(name, repo_id)))
+        response = requests.get(url, stream=True, verify=False, headers=headers)
+
+        return response.status_code == 200, url, response.status_code
+
+
     def download_model_file(self):
         try:
             while not self.download_stop and not self.file_queue.empty():

diff --git a/service/web_api.py b/service/web_api.py
@@ -195,6 +195,18 @@ def check_model_exist():
         result_list.append({"repo_id": repo_id, "type": type, "exist": exist})
     return jsonify({"code": 0, "message": "success", "exists": result_list})
 
+@app.route("/api/checkModelAccess", methods=["POST"])
+def checkModelAccess():
+    repo_id, hf_token = request.get_json()
+    downloader = HFPlaygroundDownloader(hf_token)
+    valid, url, status = downloader.is_token_valid(repo_id)
+    return jsonify(
+        {
+            "valid": valid,
+            "url": url,
+            "status": status
+        }
+    )
 
 size_cache = dict()
 lock = threading.Lock()