From 828be63f2cc100de0386b4edf77c6c597f6a14bb Mon Sep 17 00:00:00 2001 From: Anthony Guijarro Date: Sat, 27 Jan 2024 14:13:09 -0600 Subject: [PATCH] Downloader: use HF get_token function (#5381) --- download-model.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/download-model.py b/download-model.py index b22405ec..5e62036f 100644 --- a/download-model.py +++ b/download-model.py @@ -20,6 +20,7 @@ import requests import tqdm from requests.adapters import HTTPAdapter from tqdm.contrib.concurrent import thread_map +from huggingface_hub import get_token base = "https://huggingface.co" @@ -32,8 +33,8 @@ class ModelDownloader: self.session.mount('https://huggingface.co', HTTPAdapter(max_retries=max_retries)) if os.getenv('HF_USER') is not None and os.getenv('HF_PASS') is not None: self.session.auth = (os.getenv('HF_USER'), os.getenv('HF_PASS')) - if os.getenv('HF_TOKEN') is not None: - self.session.headers = {'authorization': f'Bearer {os.getenv("HF_TOKEN")}'} + if get_token() is not None: + self.session.headers = {'authorization': f'Bearer {get_token()}'} def sanitize_model_and_branch_names(self, model, branch): if model[-1] == '/':