From 5732a8f15a9af40f58c13a741b14cd72a0811d90 Mon Sep 17 00:00:00 2001 From: Gnome Ann <> Date: Wed, 11 May 2022 14:40:31 -0400 Subject: [PATCH 1/7] Don't use `aria2_hook()` if `force_download=True` is used --- utils.py | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/utils.py b/utils.py index 38066ed0..7295c677 100644 --- a/utils.py +++ b/utils.py @@ -1,4 +1,4 @@ -from threading import Timer +from threading import Timer, local import re import shutil import json @@ -141,7 +141,9 @@ def aria2_hook(pretrained_model_name_or_path: str, force_download=False, cache_d from huggingface_hub import HfFolder if shutil.which("aria2c") is None: # Don't do anything if aria2 is not installed return - if os.path.isdir(pretrained_model_name_or_path) or os.path.isfile(pretrained_model_name_or_path) or transformers.modeling_utils.is_remote_url(pretrained_model_name_or_path) or os.path.isfile(pretrained_model_name_or_path + ".index"): + if local_files_only: # If local_files_only is true, we obviously don't need to download anything + return + if os.path.isdir(pretrained_model_name_or_path) or os.path.isfile(pretrained_model_name_or_path) or os.path.isfile(pretrained_model_name_or_path + ".index") or transformers.modeling_utils.is_remote_url(pretrained_model_name_or_path): return if proxies: print("WARNING: KoboldAI does not support using aria2 to download models from huggingface.co through a proxy. Disabling aria2 download mode.") @@ -179,7 +181,7 @@ def aria2_hook(pretrained_model_name_or_path: str, force_download=False, cache_d if not sharded: # If the model has a pytorch_model.bin file, that's the only file to download filenames = [transformers.modeling_utils.WEIGHTS_NAME] else: # Otherwise download the pytorch_model.bin.index.json and then let aria2 download all the pytorch_model-#####-of-#####.bin files mentioned inside it - map_filename = transformers.file_utils.cached_path(url, cache_dir=cache_dir, force_download=force_download, proxies=proxies, resume_download=resume_download, local_files_only=local_files_only, use_auth_token=use_auth_token, user_agent=user_agent) + map_filename = transformers.file_utils.cached_path(url, cache_dir=cache_dir, force_download=force_download, proxies=proxies, resume_download=resume_download, use_auth_token=use_auth_token, user_agent=user_agent) with open(map_filename) as f: map_data = json.load(f) filenames = set(map_data["weight_map"].values()) From f60c7d84922437d4b2a614b4895b9f8cdfaed6bb Mon Sep 17 00:00:00 2001 From: Gnome Ann <> Date: Wed, 11 May 2022 14:41:34 -0400 Subject: [PATCH 2/7] Fix the behaviour of `aria2_hook()` when using `force_download` --- utils.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/utils.py b/utils.py index 7295c677..799e2e79 100644 --- a/utils.py +++ b/utils.py @@ -198,6 +198,9 @@ def aria2_hook(pretrained_model_name_or_path: str, force_download=False, cache_d path = os.path.join(_cache_dir, n + ".json") if os.path.exists(path): os.remove(path) + path = os.path.join(_cache_dir, n) + if os.path.exists(path): + os.remove(path) aria2_config = "\n".join(f"{u}\n out={n}" for u, n in zip(urls, filenames)).encode() with tempfile.NamedTemporaryFile("w+b", delete=False) as f: f.write(aria2_config) From f96c878d83bb595bb4e1d73b3ab19fcbc7695137 Mon Sep 17 00:00:00 2001 From: Gnome Ann <> Date: Wed, 11 May 2022 14:43:56 -0400 Subject: [PATCH 3/7] Use aria2 even when all model files are already in cache This allows aria2 to continue downloading a pytorch_model.bin after a cancelled download. --- utils.py | 5 ----- 1 file changed, 5 deletions(-) diff --git a/utils.py b/utils.py index 799e2e79..f8cb8027 100644 --- a/utils.py +++ b/utils.py @@ -186,11 +186,6 @@ def aria2_hook(pretrained_model_name_or_path: str, force_download=False, cache_d map_data = json.load(f) filenames = set(map_data["weight_map"].values()) urls = [transformers.file_utils.hf_bucket_url(pretrained_model_name_or_path, n, revision=revision, mirror=mirror) for n in filenames] - if not force_download: - if all(is_cached(u) for u in urls): - return - elif local_files_only: - raise FileNotFoundError("Cannot find the requested files in the cached path and outgoing traffic has been disabled. To enable model look-ups and downloads online, set 'local_files_only' to False.") etags = [h.get("X-Linked-Etag") or h.get("ETag") for u in urls for h in [requests.head(u, headers=headers, allow_redirects=False, proxies=proxies, timeout=10).headers]] filenames = [transformers.file_utils.url_to_filename(u, t) for u, t in zip(urls, etags)] if force_download: From c81f3bd0845b8e74856992d600dd104cb1868dce Mon Sep 17 00:00:00 2001 From: Gnome Ann <> Date: Wed, 11 May 2022 14:51:43 -0400 Subject: [PATCH 4/7] Use `--file-allocation=trunc` instead of `--file-allocation=none` --- utils.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/utils.py b/utils.py index f8cb8027..c541e896 100644 --- a/utils.py +++ b/utils.py @@ -200,7 +200,7 @@ def aria2_hook(pretrained_model_name_or_path: str, force_download=False, cache_d with tempfile.NamedTemporaryFile("w+b", delete=False) as f: f.write(aria2_config) f.flush() - p = subprocess.Popen(["aria2c", "-x", "10", "-s", "10", "-j", "10", "--disable-ipv6", "--file-allocation=none", "-d", _cache_dir, "-i", f.name, "-U", transformers.file_utils.http_user_agent(user_agent)] + (["-c"] if not force_download else []) + ([f"--header='Authorization: Bearer {token}'"] if use_auth_token else []), stdout=subprocess.PIPE, stderr=subprocess.STDOUT) + p = subprocess.Popen(["aria2c", "-x", "10", "-s", "10", "-j", "10", "--disable-ipv6", "--file-allocation=trunc", "-d", _cache_dir, "-i", f.name, "-U", transformers.file_utils.http_user_agent(user_agent)] + (["-c"] if not force_download else []) + ([f"--header='Authorization: Bearer {token}'"] if use_auth_token else []), stdout=subprocess.PIPE, stderr=subprocess.STDOUT) for line in p.stdout: print(line.decode(), end="", flush=True) path = f.name From 7a3f865e3f7c2fccb7951dfdf1e1cb43d052e526 Mon Sep 17 00:00:00 2001 From: Gnome Ann <> Date: Wed, 11 May 2022 15:14:37 -0400 Subject: [PATCH 5/7] Prevent aria2 from resuming cancelled downloads Resumed downloads tend to be very slow. The original transformers downloader didn't allow resuming downloads either. --- utils.py | 10 +++++++--- 1 file changed, 7 insertions(+), 3 deletions(-) diff --git a/utils.py b/utils.py index c541e896..c864a6aa 100644 --- a/utils.py +++ b/utils.py @@ -188,8 +188,12 @@ def aria2_hook(pretrained_model_name_or_path: str, force_download=False, cache_d urls = [transformers.file_utils.hf_bucket_url(pretrained_model_name_or_path, n, revision=revision, mirror=mirror) for n in filenames] etags = [h.get("X-Linked-Etag") or h.get("ETag") for u in urls for h in [requests.head(u, headers=headers, allow_redirects=False, proxies=proxies, timeout=10).headers]] filenames = [transformers.file_utils.url_to_filename(u, t) for u, t in zip(urls, etags)] - if force_download: - for n in filenames: + for n in filenames: + path = os.path.join(_cache_dir, n + ".aria2") # Prevent aria2 from continuing cancelled downloads because continued downloads are usually limited to the speed of 1 connection + if os.path.exists(path): + os.remove(path) + os.remove(os.path.join(_cache_dir, n)) + if force_download: path = os.path.join(_cache_dir, n + ".json") if os.path.exists(path): os.remove(path) @@ -200,7 +204,7 @@ def aria2_hook(pretrained_model_name_or_path: str, force_download=False, cache_d with tempfile.NamedTemporaryFile("w+b", delete=False) as f: f.write(aria2_config) f.flush() - p = subprocess.Popen(["aria2c", "-x", "10", "-s", "10", "-j", "10", "--disable-ipv6", "--file-allocation=trunc", "-d", _cache_dir, "-i", f.name, "-U", transformers.file_utils.http_user_agent(user_agent)] + (["-c"] if not force_download else []) + ([f"--header='Authorization: Bearer {token}'"] if use_auth_token else []), stdout=subprocess.PIPE, stderr=subprocess.STDOUT) + p = subprocess.Popen(["aria2c", "-x", "10", "-s", "10", "-j", "10", "--disable-ipv6", "--file-allocation=trunc", "--allow-overwrite", "--auto-file-renaming", "false", "-d", _cache_dir, "-i", f.name, "-U", transformers.file_utils.http_user_agent(user_agent)] + (["-c"] if not force_download else []) + ([f"--header='Authorization: Bearer {token}'"] if use_auth_token else []), stdout=subprocess.PIPE, stderr=subprocess.STDOUT) for line in p.stdout: print(line.decode(), end="", flush=True) path = f.name From c65272052a516e0fd617054c28a90ed16ae3a505 Mon Sep 17 00:00:00 2001 From: Gnome Ann <> Date: Wed, 11 May 2022 15:45:38 -0400 Subject: [PATCH 6/7] aria2 now downloads to different filename and renames afterwards This is to match the behaviour of the original transformers downloader in order to deal with the rare case of someone downloading a model using aria2, cancelling before it finishes, and then attempting to resume the download with the normal transformers downloader. --- utils.py | 15 +++++++++++---- 1 file changed, 11 insertions(+), 4 deletions(-) diff --git a/utils.py b/utils.py index c864a6aa..f65222a4 100644 --- a/utils.py +++ b/utils.py @@ -1,4 +1,4 @@ -from threading import Timer, local +from threading import Timer import re import shutil import json @@ -186,13 +186,19 @@ def aria2_hook(pretrained_model_name_or_path: str, force_download=False, cache_d map_data = json.load(f) filenames = set(map_data["weight_map"].values()) urls = [transformers.file_utils.hf_bucket_url(pretrained_model_name_or_path, n, revision=revision, mirror=mirror) for n in filenames] + if not force_download: + urls = [u for u in urls if not is_cached(u)] + if not urls: + return etags = [h.get("X-Linked-Etag") or h.get("ETag") for u in urls for h in [requests.head(u, headers=headers, allow_redirects=False, proxies=proxies, timeout=10).headers]] filenames = [transformers.file_utils.url_to_filename(u, t) for u, t in zip(urls, etags)] for n in filenames: - path = os.path.join(_cache_dir, n + ".aria2") # Prevent aria2 from continuing cancelled downloads because continued downloads are usually limited to the speed of 1 connection + path = os.path.join(_cache_dir, "kai-tempfile." + n + ".aria2") + if os.path.exists(path): + os.remove(path) + path = os.path.join(_cache_dir, "kai-tempfile." + n) if os.path.exists(path): os.remove(path) - os.remove(os.path.join(_cache_dir, n)) if force_download: path = os.path.join(_cache_dir, n + ".json") if os.path.exists(path): @@ -200,7 +206,7 @@ def aria2_hook(pretrained_model_name_or_path: str, force_download=False, cache_d path = os.path.join(_cache_dir, n) if os.path.exists(path): os.remove(path) - aria2_config = "\n".join(f"{u}\n out={n}" for u, n in zip(urls, filenames)).encode() + aria2_config = "\n".join(f"{u}\n out=kai-tempfile.{n}" for u, n in zip(urls, filenames)).encode() with tempfile.NamedTemporaryFile("w+b", delete=False) as f: f.write(aria2_config) f.flush() @@ -213,5 +219,6 @@ def aria2_hook(pretrained_model_name_or_path: str, force_download=False, cache_d except OSError: pass for u, t, n in zip(urls, etags, filenames): + os.rename(os.path.join(_cache_dir, "kai-tempfile." + n), os.path.join(_cache_dir, n)) with open(os.path.join(_cache_dir, n + ".json"), "w") as f: json.dump({"url": u, "etag": t}, f) From 2ebba9488b3096335a9beaa6e806552f93dfa211 Mon Sep 17 00:00:00 2001 From: Gnome Ann <> Date: Wed, 11 May 2022 15:51:48 -0400 Subject: [PATCH 7/7] Change `force_download` back to False This is to prevent fully downloaded models from being re-downloaded in Colab. --- utils.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/utils.py b/utils.py index 12876f6f..f65222a4 100644 --- a/utils.py +++ b/utils.py @@ -135,7 +135,7 @@ def decodenewlines(txt): #==================================================================# # Downloads sharded huggingface checkpoints using aria2c if possible #==================================================================# -def aria2_hook(pretrained_model_name_or_path: str, force_download=True, cache_dir=None, proxies=None, resume_download=False, local_files_only=False, use_auth_token=None, user_agent=None, revision=None, mirror=None, **kwargs): +def aria2_hook(pretrained_model_name_or_path: str, force_download=False, cache_dir=None, proxies=None, resume_download=False, local_files_only=False, use_auth_token=None, user_agent=None, revision=None, mirror=None, **kwargs): import transformers import transformers.modeling_utils from huggingface_hub import HfFolder