rixty/KoboldAI-Client
1
1
Fork 0
mirror of https://github.com/KoboldAI/KoboldAI-Client.git synced 2025-06-05 21:59:24 +02:00
Code Issues Projects Releases Wiki Activity

Initial World Info showing in new ui

Browse Source 
Added model download status to old ui
This commit is contained in:
ebolam
2022-07-22 14:05:21 -04:00
parent fe76d64f92
commit aff889bb69
4 changed files with 324 additions and 82 deletions
Download Patch File Download Diff File Expand all files Collapse all files

12
utils.py
View File

@@ -169,6 +169,16 @@ def num_layers(config):
#==================================================================#
# Downloads huggingface checkpoints using aria2c if possible
#==================================================================#
from flask_socketio import emit
class Send_to_socketio(object):
def write(self, bar):
print("should be emitting: ", bar, end="")
time.sleep(0.01)
try:
emit('from_server', {'cmd': 'model_load_status', 'data': bar.replace(" ", " ")}, broadcast=True, room="UI_1")
except:
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):
import transformers
import transformers.modeling_utils
@@ -265,7 +275,7 @@ def aria2_hook(pretrained_model_name_or_path: str, force_download=False, cache_d
done = True
break
if bar is None:
bar = tqdm(total=total_length, desc=f"[aria2] Downloading model", unit="B", unit_scale=True, unit_divisor=1000)
bar = tqdm(total=total_length, desc=f"[aria2] Downloading model", unit="B", unit_scale=True, unit_divisor=1000, file=Send_to_socketio())
visited = set()
for x in r:
filename = x["files"][0]["path"]