From 04862034c9f043a1453e46779036d1a9ca325f85 Mon Sep 17 00:00:00 2001
From: zanllp <qc@zanllp.cn>
Date: Mon, 5 Jan 2026 06:42:21 +0800
Subject: [PATCH] Refactor install script: simplify package installation logic
 and improve error handling

---
 install.py | 40 +++++++++++++++++++++++++++++-----------
 1 file changed, 29 insertions(+), 11 deletions(-)

diff --git a/install.py b/install.py
index a0acdf9..017d08c 100644
--- a/install.py
+++ b/install.py
@@ -1,6 +1,5 @@
 import launch
 import os
-import pkg_resources
 
 req_file = os.path.join(os.path.dirname(os.path.realpath(__file__)), "requirements.txt")
 
@@ -11,18 +10,37 @@ def dist2package(dist: str):
         "pillow-avif-plugin": "pillow_avif"
     }).get(dist, dist)
 
-# copy from controlnet, thanks
+def extract_package_name(package_spec: str) -> str:
+    """Extract package name from package specification.
+    Handles formats like: package, package==1.0, package>=1.0, package>=1.0,<2.0
+    """
+    # Remove git URL suffix if present
+    package_spec = package_spec.split("@git")[0].strip()
+    # Extract package name (everything before version specifiers)
+    for sep in ["==", ">=", "<=", ">", "<", "~=", "!="]:
+        if sep in package_spec:
+            package_spec = package_spec.split(sep)[0].strip()
+            break
+    # Handle version ranges - extract package name before first comma
+    if "," in package_spec:
+        package_spec = package_spec.split(",")[0].strip()
+    return package_spec
+
+
+
 with open(req_file) as file:
     for package in file:
         try:
             package = package.strip()
-            if '==' in package:
-                package_name, package_version = package.split('==')
-                installed_version = pkg_resources.get_distribution(package_name).version
-                if installed_version != package_version:
-                    launch.run_pip(f"install {package}", f"sd-webui-infinite-image-browsing requirement: changing {package_name} version from {installed_version} to {package_version}")
-            elif not launch.is_installed(dist2package(package)):
-                launch.run_pip(f"install {package}", f"sd-webui-infinite-image-browsing requirement: {package}")
+            package_name = extract_package_name(package)
+            if not launch.is_installed(dist2package(package_name)):
+                launch.run_pip(f"install \"{package}\"", f"sd-webui-infinite-image-browsing requirement: {package}")
         except Exception as e:
-            print(e)
-            print(f'Warning: Failed to install {package}, something may not work.')
+            # Handle Unicode encoding errors on Windows
+            try:
+                print(str(e))
+            except UnicodeEncodeError:
+                print(repr(e))
+            print(
+                f"Warning: Failed to install {package}, some preprocessors may not work."
+            )