AI/text-generation-webui
1
0
Fork 0
mirror of https://github.com/oobabooga/text-generation-webui.git synced 2024-10-01 01:26:03 -04:00
Code Issues Packages Projects Releases Wiki Activity

downloader: use 1 session for all files for better speed

Browse Source
This commit is contained in:
oobabooga 2024-08-06 19:41:12 -07:00
parent 608545d282
commit f106e780ba
1 changed files with 3 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

5
download-model.py
View File

@ -29,6 +29,7 @@ base = os.environ.get("HF_ENDPOINT") or "https://huggingface.co"
class ModelDownloader: class ModelDownloader:
def __init__(self, max_retries=5): def __init__(self, max_retries=5):
self.max_retries = max_retries self.max_retries = max_retries
self.session = self.get_session()
def get_session(self): def get_session(self):
session = requests.Session() session = requests.Session()
@ -72,7 +73,7 @@ class ModelDownloader:
return model, branch return model, branch
def get_download_links_from_huggingface(self, model, branch, text_only=False, specific_file=None): def get_download_links_from_huggingface(self, model, branch, text_only=False, specific_file=None):
session = self.get_session() session = self.session
page = f"/api/models/{model}/tree/{branch}" page = f"/api/models/{model}/tree/{branch}"
cursor = b"" cursor = b""
@ -192,7 +193,7 @@ class ModelDownloader:
attempt = 0 attempt = 0
while attempt < max_retries: while attempt < max_retries:
attempt += 1 attempt += 1
session = self.get_session() session = self.session
headers = {} headers = {}
mode = 'wb' mode = 'wb'