Fix for getting "model download status" when downloading config to figure out layer counts

This commit is contained in:
ebolam 2022-07-25 18:29:14 -04:00
parent 907cf74b13
commit 12acb50ee0
1 changed files with 15 additions and 11 deletions

View File

@ -1312,6 +1312,7 @@ def patch_transformers_download():
def write(self, bar): def write(self, bar):
bar = bar.replace("\r", "") bar = bar.replace("\r", "")
try: try:
print(bar, end="\r")
emit('from_server', {'cmd': 'model_load_status', 'data': bar.replace(" ", " ")}, broadcast=True) emit('from_server', {'cmd': 'model_load_status', 'data': bar.replace(" ", " ")}, broadcast=True)
eventlet.sleep(seconds=0) eventlet.sleep(seconds=0)
except: except:
@ -1336,6 +1337,7 @@ def patch_transformers_download():
total = resume_size + int(content_length) if content_length is not None else None total = resume_size + int(content_length) if content_length is not None else None
# `tqdm` behavior is determined by `utils.logging.is_progress_bar_enabled()` # `tqdm` behavior is determined by `utils.logging.is_progress_bar_enabled()`
# and can be set using `utils.logging.enable/disable_progress_bar()` # and can be set using `utils.logging.enable/disable_progress_bar()`
if url[-11:] != 'config.json':
progress = tqdm.tqdm( progress = tqdm.tqdm(
unit="B", unit="B",
unit_scale=True, unit_scale=True,
@ -1347,8 +1349,10 @@ def patch_transformers_download():
) )
for chunk in r.iter_content(chunk_size=1024): for chunk in r.iter_content(chunk_size=1024):
if chunk: # filter out keep-alive new chunks if chunk: # filter out keep-alive new chunks
if url[-11:] != 'config.json':
progress.update(len(chunk)) progress.update(len(chunk))
temp_file.write(chunk) temp_file.write(chunk)
if url[-11:] != 'config.json':
progress.close() progress.close()
transformers.utils.hub.http_get = http_get transformers.utils.hub.http_get = http_get