From 0e270e0b2587482883fe66dadd8296828e72ee03 Mon Sep 17 00:00:00 2001
From: ebolam <ebolam@gmail.com>
Date: Wed, 21 Sep 2022 13:07:49 -0400
Subject: [PATCH] Fix for aria2 download status to UI
---
utils.py | 16 +++++++++++++++-
1 file changed, 15 insertions(+), 1 deletion(-)
diff --git a/utils.py b/utils.py
index 34562aa5..696a6d87 100644
--- a/utils.py
+++ b/utils.py
@@ -188,6 +188,20 @@ class Send_to_socketio(object):
pass
def _download_with_aria2(aria2_config: str, total_length: int, directory: str = ".", user_agent=None, force_download=False, use_auth_token=None):
+ class Send_to_socketio(object):
+ def write(self, bar):
+ bar = bar.replace("\r", "").replace("\n", "")
+
+ if bar != "":
+ try:
+ print(bar, end="\n")
+ emit('from_server', {'cmd': 'model_load_status', 'data': bar.replace(" ", " ")}, broadcast=True)
+ eventlet.sleep(seconds=0)
+ except:
+ pass
+ def flush(self):
+ pass
+
import transformers
lengths = {}
s = requests.Session()
@@ -211,7 +225,7 @@ def _download_with_aria2(aria2_config: str, total_length: int, directory: str =
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"]