From 92ce2d706b5c29cd0646a26d765a7d4a9a98b1c6 Mon Sep 17 00:00:00 2001
From: "Dr.Lt.Data" <dr.lt.data@gmail.com>
Date: Sun, 5 Nov 2023 10:05:39 +0900
Subject: [PATCH] improve: tqdm in process fix: encoding error of process
 output

---
 __init__.py          | 12 ++++++------
 prestartup_script.py | 43 +++++++++++++++++++++++++++++++------------
 2 files changed, 37 insertions(+), 18 deletions(-)

diff --git a/__init__.py b/__init__.py
index cd244146..f50c4620 100644
--- a/__init__.py
+++ b/__init__.py
@@ -6,15 +6,15 @@ import sys
 import threading
 import datetime
 import re
+import locale
 from tqdm.auto import tqdm
 
+version = "V0.39"
 
 def handle_stream(stream, prefix):
-    try:
-        for line in stream:
-            print(prefix, line, end="")
-    except Exception:
-        print("[!] ??? log decoding error ???")
+    stream.reconfigure(encoding=locale.getpreferredencoding(), errors='replace')
+    for line in stream:
+        print(prefix, line, end="")
 
 
 def run_script(cmd, cwd='.'):
@@ -63,7 +63,7 @@ sys.path.append('../..')
 from torchvision.datasets.utils import download_url
 
 # ensure .js
-print("### Loading: ComfyUI-Manager (V0.38.3)")
+print(f"### Loading: ComfyUI-Manager ({version})")
 
 comfy_ui_required_revision = 1240
 comfy_ui_revision = "Unknown"
diff --git a/prestartup_script.py b/prestartup_script.py
index f80c1a36..65053133 100644
--- a/prestartup_script.py
+++ b/prestartup_script.py
@@ -5,6 +5,7 @@ import sys
 import atexit
 import threading
 import re
+import locale
 
 
 message_collapses = []
@@ -25,11 +26,18 @@ git_script_path = os.path.join(comfyui_manager_path, "git_helper.py")
 
 
 def handle_stream(stream, prefix):
-    try:
-        for msg in stream:
-            print(prefix, msg, end="")
-    except Exception:
-        print("[!] ??? log decoding error ???")
+    stream.reconfigure(encoding=locale.getpreferredencoding(), errors='replace')
+    for msg in stream:
+        if prefix == '[!]' and ('it/s]' or 's/it]') in msg and ('%|' in msg or 'it [' in msg):
+            if msg.startswith('100%'):
+                print('\r' + msg, end="", file=sys.stderr),
+            else:
+                print('\r' + msg[:-1], end="", file=sys.stderr),
+        else:
+            if prefix == '[!]':
+                print(prefix, msg, end="", file=sys.stderr)
+            else:
+                print(prefix, msg, end="")
 
 
 def process_wrap(cmd_str, cwd_path, handler=None):
@@ -150,14 +158,25 @@ if os.path.exists(restore_snapshot_path):
         cloned_repos = []
 
         def msg_capture(stream, prefix):
-            try:
-                for msg in stream:
-                    if msg.startswith("CLONE: "):
-                        cloned_repos.append(msg[7:])
+            stream.reconfigure(encoding=locale.getpreferredencoding(), errors='replace')
+            for msg in stream:
+                if msg.startswith("CLONE: "):
+                    cloned_repos.append(msg[7:])
+                    if prefix == '[!]':
+                        print(prefix, msg, end="", file=sys.stderr)
+                    else:
+                        print(prefix, msg, end="")
 
-                    print(prefix, msg, end="")
-            except Exception:
-                print("[!] [snapshot restore] ??? log decoding error ???")
+                elif prefix == '[!]' and ('it/s]' in msg or 's/it]' in msg) and ('%|' in msg or 'it [' in msg):
+                    if msg.startswith('100%'):
+                        print('\r' + msg, end="", file=sys.stderr),
+                    else:
+                        print('\r'+msg[:-1], end="", file=sys.stderr),
+                else:
+                    if prefix == '[!]':
+                        print(prefix, msg, end="", file=sys.stderr)
+                    else:
+                        print(prefix, msg, end="")
 
         print(f"[ComfyUI-Manager] Restore snapshot.")
         cmd_str = [sys.executable, git_script_path, '--apply-snapshot', restore_snapshot_path]