Merge pull request #16869 from Haoming02/uv-hook

Support uv Package Manager

launch.py

@@ -6,6 +6,11 @@ git = launch_utils.git
 index_url = launch_utils.index_url
 dir_repos = launch_utils.dir_repos
 
+if args.uv:
+    from modules.uv_hook import patch
+    patch()
+
+
 commit_hash = launch_utils.commit_hash
 git_tag = launch_utils.git_tag
 

modules/cmd_args.py

@@ -126,3 +126,4 @@ parser.add_argument("--skip-load-model-at-start", action='store_true', help="if
 parser.add_argument("--unix-filenames-sanitization", action='store_true', help="allow any symbols except '/' in filenames. May conflict with your browser and file system")
 parser.add_argument("--filenames-max-length", type=int, default=128, help='maximal length of filenames of saved images. If you override it, it can conflict with your file system')
 parser.add_argument("--no-prompt-history", action='store_true', help="disable read prompt from last generation feature; settings this argument will not create '--data_path/params.txt' file")
+parser.add_argument("--uv", action='store_true', help="use the uv package manager")

modules/uv_hook.py

@@ -0,0 +1,49 @@
+import sys
+import copy
+import shlex
+import subprocess
+from functools import wraps
+
+BAD_FLAGS = ("--prefer-binary", '-I', '--ignore-installed')
+
+
+def patch():
+    if hasattr(subprocess, "__original_run"):
+        return
+
+    print("using uv")
+    try:
+        subprocess.run(['uv', '-V'])
+    except FileNotFoundError:
+        subprocess.run([sys.executable, '-m', 'pip', 'install', 'uv'])
+
+    subprocess.__original_run = subprocess.run
+
+    @wraps(subprocess.__original_run)
+    def patched_run(*args, **kwargs):
+        _kwargs = copy.copy(kwargs)
+        if args:
+            command, *_args = args
+        else:
+            command, _args = _kwargs.pop("args", ""), ()
+
+        if isinstance(command, str):
+            command = shlex.split(command)
+        else:
+            command = [arg.strip() for arg in command]
+
+        if not isinstance(command, list) or "pip" not in command:
+            return subprocess.__original_run(*args, **kwargs)
+
+        cmd = command[command.index("pip") + 1:]
+
+        cmd = [arg for arg in cmd if arg not in BAD_FLAGS]
+
+        modified_command = ["uv", "pip", *cmd]
+
+        result = subprocess.__original_run([*modified_command, *_args], **_kwargs)
+        if result.returncode != 0:
+            return subprocess.__original_run(*args, **kwargs)
+        return result
+
+    subprocess.run = patched_run