From f106e780ba91ad8c2949596f506a7f762e3724c5 Mon Sep 17 00:00:00 2001
From: oobabooga <112222186+oobabooga@users.noreply.github.com>
Date: Tue, 6 Aug 2024 19:41:12 -0700
Subject: [PATCH] downloader: use 1 session for all files for better speed

---
 download-model.py | 5 +++--
 1 file changed, 3 insertions(+), 2 deletions(-)

diff --git a/download-model.py b/download-model.py
index 0014b689..306784a3 100644
--- a/download-model.py
+++ b/download-model.py
@@ -29,6 +29,7 @@ base = os.environ.get("HF_ENDPOINT") or "https://huggingface.co"
 class ModelDownloader:
     def __init__(self, max_retries=5):
         self.max_retries = max_retries
+        self.session = self.get_session()
 
     def get_session(self):
         session = requests.Session()
@@ -72,7 +73,7 @@ class ModelDownloader:
         return model, branch
 
     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}"
         cursor = b""
 
@@ -192,7 +193,7 @@ class ModelDownloader:
         attempt = 0
         while attempt < max_retries:
             attempt += 1
-            session = self.get_session()
+            session = self.session
             headers = {}
             mode = 'wb'