Fix for aria2 download status to UI
This commit is contained in:
parent
f3d6beb578
commit
8915ee7eb3
17
utils.py
17
utils.py
|
@ -178,16 +178,25 @@ def num_layers(config):
|
||||||
# Downloads huggingface checkpoints using aria2c if possible
|
# Downloads huggingface checkpoints using aria2c if possible
|
||||||
#==================================================================#
|
#==================================================================#
|
||||||
from flask_socketio import emit
|
from flask_socketio import emit
|
||||||
|
|
||||||
|
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):
|
class Send_to_socketio(object):
|
||||||
def write(self, bar):
|
def write(self, bar):
|
||||||
time.sleep(0.01)
|
bar = bar.replace("\r", "").replace("\n", "")
|
||||||
|
|
||||||
|
if bar != "":
|
||||||
|
try:
|
||||||
|
print(bar, end="\n")
|
||||||
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)
|
||||||
except:
|
except:
|
||||||
pass
|
pass
|
||||||
|
eventlet.sleep(seconds=0)
|
||||||
|
except:
|
||||||
|
pass
|
||||||
|
def flush(self):
|
||||||
|
pass
|
||||||
|
|
||||||
def _download_with_aria2(aria2_config: str, total_length: int, directory: str = ".", user_agent=None, force_download=False, use_auth_token=None):
|
|
||||||
import transformers
|
import transformers
|
||||||
lengths = {}
|
lengths = {}
|
||||||
s = requests.Session()
|
s = requests.Session()
|
||||||
|
@ -211,7 +220,7 @@ def _download_with_aria2(aria2_config: str, total_length: int, directory: str =
|
||||||
done = True
|
done = True
|
||||||
break
|
break
|
||||||
if bar is None:
|
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()
|
visited = set()
|
||||||
for x in r:
|
for x in r:
|
||||||
filename = x["files"][0]["path"]
|
filename = x["files"][0]["path"]
|
||||||
|
|
Loading…
Reference in New Issue