Fix for Colab download status bar

This commit is contained in:
ebolam 2022-08-19 12:05:22 -04:00
parent cf422aa16e
commit 046f9d8ace
1 changed files with 2 additions and 0 deletions

View File

@ -181,6 +181,7 @@ class Send_to_socketio(object):
print(bar, end="\r") 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)
except: except:
raise
pass pass
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): 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):
@ -279,6 +280,7 @@ def aria2_hook(pretrained_model_name_or_path: str, force_download=False, cache_d
done = True done = True
break break
if bar is None: if bar is None:
print("setting up status bar for aria2 download")
bar = tqdm(total=total_length, desc=f"[aria2] Downloading model", unit="B", unit_scale=True, unit_divisor=1000, file=Send_to_socketio()) bar = tqdm(total=total_length, desc=f"[aria2] Downloading model", unit="B", unit_scale=True, unit_divisor=1000, file=Send_to_socketio())
visited = set() visited = set()
for x in r: for x in r: