Merge pull request #214 from ebolam/united

Fix for Aria2 Status in UI

utils.py

@@ -178,16 +178,25 @@ 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):
-        time.sleep(0.01)
-        try:
-            print(bar, end = "\r")
-            emit('from_server', {'cmd': 'model_load_status', 'data': bar.replace(" ", " ")}, broadcast=True)
-        except:
-            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")
+                    try:
+                        emit('from_server', {'cmd': 'model_load_status', 'data': bar.replace(" ", " ")}, broadcast=True)
+                    except:
+                        pass
+                    eventlet.sleep(seconds=0)
+                except:
+                    pass
+        def flush(self):
+            pass
+    
     import transformers
     lengths = {}
     s = requests.Session()
@@ -211,7 +220,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"]