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

Add a timeout to download-model.py requests

Browse Source
This commit is contained in:
oobabooga 2023-05-19 11:19:34 -03:00
parent f052ab9c8f
commit 39dab18307
1 changed files with 3 additions and 3 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

6
download-model.py
View File

@ -98,7 +98,7 @@ def get_download_links_from_huggingface(model, branch, text_only=False):
is_lora = False is_lora = False
while True: while True:
url = f"{base}{page}" + (f"?cursor={cursor.decode()}" if cursor else "") url = f"{base}{page}" + (f"?cursor={cursor.decode()}" if cursor else "")
r = requests.get(url) r = requests.get(url, timeout=10)
r.raise_for_status() r.raise_for_status()
content = r.content content = r.content
@ -169,7 +169,7 @@ def get_single_file(url, output_folder, start_from_scratch=False):
output_path = output_folder / filename output_path = output_folder / filename
if output_path.exists() and not start_from_scratch: if output_path.exists() and not start_from_scratch:
# Check if the file has already been downloaded completely # Check if the file has already been downloaded completely
r = requests.get(url, stream=True) r = requests.get(url, stream=True, timeout=10)
total_size = int(r.headers.get('content-length', 0)) total_size = int(r.headers.get('content-length', 0))
if output_path.stat().st_size >= total_size: if output_path.stat().st_size >= total_size:
return return
@ -180,7 +180,7 @@ def get_single_file(url, output_folder, start_from_scratch=False):
headers = {} headers = {}
mode = 'wb' mode = 'wb'
r = requests.get(url, stream=True, headers=headers) r = requests.get(url, stream=True, headers=headers, timeout=10)
with open(output_path, mode) as f: with open(output_path, mode) as f:
total_size = int(r.headers.get('content-length', 0)) total_size = int(r.headers.get('content-length', 0))
block_size = 1024 block_size = 1024