From e497b8321f2face82d3f1491a4c3e05834703d76 Mon Sep 17 00:00:00 2001
From: "Dr.Lt.Data" <dr.lt.data@gmail.com>
Date: Sat, 4 Nov 2023 21:23:29 +0900
Subject: [PATCH] feat: tqdm for clone fix: snapshot restore proper working
 directory

---
 __init__.py          | 19 ++++++++++++++++---
 git_helper.py        | 16 +++++++++++++++-
 prestartup_script.py |  2 ++
 3 files changed, 33 insertions(+), 4 deletions(-)

diff --git a/__init__.py b/__init__.py
index 89b8e01a..df8d8e19 100644
--- a/__init__.py
+++ b/__init__.py
@@ -6,7 +6,8 @@ import sys
 import threading
 import datetime
 import re
-
+from tqdm.auto import tqdm
+from git.remote import RemoteProgress
 
 def handle_stream(stream, prefix):
     for line in stream:
@@ -57,7 +58,7 @@ sys.path.append('../..')
 from torchvision.datasets.utils import download_url
 
 # ensure .js
-print("### Loading: ComfyUI-Manager (V0.37)")
+print("### Loading: ComfyUI-Manager (V0.38)")
 
 comfy_ui_required_revision = 1240
 comfy_ui_revision = "Unknown"
@@ -952,6 +953,18 @@ def execute_install_script(url, repo_path):
     return True
 
 
+class GitProgress(RemoteProgress):
+    def __init__(self):
+        super().__init__()
+        self.pbar = tqdm()
+
+    def update(self, op_code, cur_count, max_count=None, message=''):
+        self.pbar.total = max_count
+        self.pbar.n = cur_count
+        self.pbar.pos = 0
+        self.pbar.refresh()
+
+
 def gitclone_install(files):
     print(f"install: {files}")
     for url in files:
@@ -966,7 +979,7 @@ def gitclone_install(files):
             if platform.system() == 'Windows':
                 run_script([sys.executable, git_script_path, "--clone", custom_nodes_path, url])
             else:
-                repo = git.Repo.clone_from(url, repo_path, recursive=True)
+                repo = git.Repo.clone_from(url, repo_path, recursive=True, progress=GitProgress())
                 repo.git.clear_cache()
                 repo.close()
 
diff --git a/git_helper.py b/git_helper.py
index 2fbcf078..73bf9c81 100644
--- a/git_helper.py
+++ b/git_helper.py
@@ -5,18 +5,32 @@ import configparser
 import re
 import json
 from torchvision.datasets.utils import download_url
+from tqdm.auto import tqdm
+from git.remote import RemoteProgress
 
 config_path = os.path.join(os.path.dirname(__file__), "config.ini")
 nodelist_path = os.path.join(os.path.dirname(__file__), "custom-node-list.json")
 working_directory = os.getcwd()
 
 
+class GitProgress(RemoteProgress):
+    def __init__(self):
+        super().__init__()
+        self.pbar = tqdm()
+
+    def update(self, op_code, cur_count, max_count=None, message=''):
+        self.pbar.total = max_count
+        self.pbar.n = cur_count
+        self.pbar.pos = 0
+        self.pbar.refresh()
+
+
 def gitclone(custom_nodes_path, url, target_hash=None):
     repo_name = os.path.splitext(os.path.basename(url))[0]
     repo_path = os.path.join(custom_nodes_path, repo_name)
 
     # Clone the repository from the remote URL
-    repo = git.Repo.clone_from(url, repo_path, recursive=True)
+    repo = git.Repo.clone_from(url, repo_path, recursive=True, progress=GitProgress())
 
     if target_hash is not None:
         print(f"CHECKOUT: {repo_name} [{target_hash}]")
diff --git a/prestartup_script.py b/prestartup_script.py
index 374e62ba..b7a63cc0 100644
--- a/prestartup_script.py
+++ b/prestartup_script.py
@@ -163,6 +163,7 @@ if os.path.exists(restore_snapshot_path):
                 try:
                     repository_name = url.split("/")[-1].strip()
                     repo_path = os.path.join(custom_nodes_path, repository_name)
+                    repo_path = os.path.abspath(repo_path)
 
                     requirements_path = os.path.join(repo_path, 'requirements.txt')
                     install_script_path = os.path.join(repo_path, 'install.py')
@@ -179,6 +180,7 @@ if os.path.exists(restore_snapshot_path):
 
                     if os.path.exists(install_script_path):
                         install_cmd = [sys.executable, install_script_path]
+                        print(f">>> {install_cmd} / {repo_path}")
                         this_exit_code += process_wrap(install_cmd, repo_path)
 
                     if this_exit_code != 0: