From d82922097388f2d618b4a031da021276aab1474c Mon Sep 17 00:00:00 2001
From: "dr.lt.data" <Dr.Lt.Data@gmail.com>
Date: Thu, 14 Mar 2024 12:45:18 +0900
Subject: [PATCH] improve: pip handling - prevent downgrade `torch, torchsde,
 torchvision, transformers, safetensors, kornia`

---
 __init__.py          | 33 ++++++++++++++++++++++++++++-----
 prestartup_script.py | 10 +++++++++-
 2 files changed, 37 insertions(+), 6 deletions(-)

diff --git a/__init__.py b/__init__.py
index fe13d91f..a6de5bed 100644
--- a/__init__.py
+++ b/__init__.py
@@ -29,7 +29,7 @@ except:
     print(f"[WARN] ComfyUI-Manager: Your ComfyUI version is outdated. Please update to the latest version.")
 
 
-version = [2, 9]
+version = [2, 10]
 version_str = f"V{version[0]}.{version[1]}" + (f'.{version[2]}' if len(version) > 2 else '')
 print(f"### Loading: ComfyUI-Manager ({version_str})")
 
@@ -39,6 +39,22 @@ comfy_ui_hash = "-"
 cache_lock = threading.Lock()
 
 
+def is_blacklisted(name):
+    name = name.strip()
+
+    pattern = r'([^<>!=]+)([<>!=]=?)'
+    match = re.search(pattern, name)
+
+    if match:
+        name = match.group(1)
+
+    if name in cm_global.pip_downgrade_blacklist:
+        if match is None or match.group(2) in ['<=', '==', '<']:
+            return True
+
+    return False
+
+
 def handle_stream(stream, prefix):
     stream.reconfigure(encoding=locale.getpreferredencoding(), errors='replace')
     for msg in stream:
@@ -286,6 +302,12 @@ def try_install_script(url, repo_path, install_cmd):
 
         return True
     else:
+        if len(install_cmd) == 5 and install_cmd[2:4] == ['pip', 'install']:
+            if is_blacklisted(install_cmd[4]):
+                print(f"[ComfyUI-Manager] skip black listed pip installation: '{install_cmd[4]}'")
+                return True
+
+
         print(f"\n## ComfyUI-Manager: EXECUTE => {install_cmd}")
         code = run_script(install_cmd, cwd=repo_path)
 
@@ -525,16 +547,17 @@ def git_pull(path):
     else:
         repo = git.Repo(path)
 
-        print(f"path={path} / repo.is_dirty: {repo.is_dirty()}")
-
         if repo.is_dirty():
             repo.git.stash()
 
         if repo.head.is_detached:
             switch_to_default_branch(repo)
 
-        origin = repo.remote(name='origin')
-        origin.pull()
+        current_branch = repo.active_branch
+        remote_name = current_branch.tracking_branch().remote_name
+        remote = repo.remote(name=remote_name)
+
+        remote.pull()
         repo.git.submodule('update', '--init', '--recursive')
         
         repo.close()
diff --git a/prestartup_script.py b/prestartup_script.py
index 663c65eb..b7e13f42 100644
--- a/prestartup_script.py
+++ b/prestartup_script.py
@@ -15,6 +15,9 @@ sys.path.append(glob_path)
 import cm_global
 
 
+cm_global.pip_downgrade_blacklist = ['torch', 'torchsde', 'torchvision', 'transformers', 'safetensors', 'kornia']
+
+
 def skip_pip_spam(x):
     return 'Requirement already satisfied:' in x
 
@@ -350,7 +353,12 @@ def is_installed(name):
     
     if match:
         name = match.group(1)
-        
+
+    if name in cm_global.pip_downgrade_blacklist:
+        if match is None or match.group(2) in ['<=', '==', '<']:
+            print(f"[ComfyUI-Manager] skip black listed pip installation: '{name}'")
+            return True
+
     return name.lower() in get_installed_packages()