diff --git a/launch.py b/launch.py index f83820d25..3c9659914 100644 --- a/launch.py +++ b/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 diff --git a/modules/cmd_args.py b/modules/cmd_args.py index d71982b2c..2ebe267cc 100644 --- a/modules/cmd_args.py +++ b/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") diff --git a/modules/uv_hook.py b/modules/uv_hook.py new file mode 100644 index 000000000..4f1392f4a --- /dev/null +++ b/modules/uv_hook.py @@ -0,0 +1,39 @@ +import shlex +import subprocess +from functools import wraps + + +def patch(): + if hasattr(subprocess, "__original_run"): + return + + print("using uv") + subprocess.__original_run = subprocess.run + + @wraps(subprocess.__original_run) + def patched_run(*args, **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] + + assert isinstance(command, list) + + if "pip" not in command: + return subprocess.__original_run([*command, *_args], **kwargs) + + cmd = command[command.index("pip") + 1 :] + + BAD_FLAGS = ("--prefer-binary",) + cmd = [arg for arg in cmd if arg not in BAD_FLAGS] + + modified_command = ["uv", "pip", *cmd] + + return subprocess.__original_run([*modified_command, *_args], **kwargs) + + subprocess.run = patched_run