From 74e61a1bce912fe9719efeb9e03415e11832b4a6 Mon Sep 17 00:00:00 2001 From: Mister K <678459+kairin@users.noreply.github.com> Date: Fri, 2 May 2025 23:48:16 +0800 Subject: [PATCH] refactor: enhance dependency update script with improved UV installation checks and error handling --- update_dependencies_uv.py | 157 ++++++++++++++++++++------------------ 1 file changed, 82 insertions(+), 75 deletions(-) diff --git a/update_dependencies_uv.py b/update_dependencies_uv.py index d488bc66b..ebd2a2101 100755 --- a/update_dependencies_uv.py +++ b/update_dependencies_uv.py @@ -1,103 +1,110 @@ #!/usr/bin/env python3 """ -ComfyUI Dependency Update Script using UV -Created: May 2, 2025 +Update dependencies for ComfyUI using UV """ -import argparse import os -import platform -import subprocess import sys +import subprocess +import argparse +from pathlib import Path def check_uv_installed(): - """Check if UV is installed and install if missing.""" + """Check if UV is installed and install it if not.""" try: subprocess.run(["uv", "--version"], stdout=subprocess.PIPE, stderr=subprocess.PIPE, check=True) - print("✅ UV is installed") - return True - except (subprocess.CalledProcessError, FileNotFoundError): - print("⚠️ UV not found. Installing...") - - if platform.system() == "Windows": - subprocess.run( - ["powershell", "-Command", "Invoke-WebRequest -Uri https://astral.sh/uv/install.ps1 -OutFile install.ps1; ./install.ps1"], - check=True - ) - else: - subprocess.run( - ["bash", "-c", "curl -LsSf https://astral.sh/uv/install.sh | sh"], - check=True - ) - - print("✅ UV installed successfully") return True + except (subprocess.SubprocessError, FileNotFoundError): + print("UV not found. Installing UV...") + try: + if sys.platform == "win32": + subprocess.run( + ["powershell", "-Command", "Invoke-WebRequest -Uri https://astral.sh/uv/install.ps1 -OutFile install.ps1; ./install.ps1"], + check=True + ) + else: # macOS or Linux + subprocess.run( + ["bash", "-c", "curl -LsSf https://astral.sh/uv/install.sh | sh"], + check=True + ) + + # Update PATH for this session + if sys.platform != "win32": + os.environ["PATH"] = f"{os.path.expanduser('~/.local/bin')}:{os.environ['PATH']}" + + print("UV installed successfully.") + return True + except subprocess.SubprocessError: + print("Failed to install UV. Please install manually from https://github.com/astral-sh/uv") + return False -def update_dependencies(args): +def update_dependencies(advanced=False, cuda=False, lock=False): """Update dependencies using UV.""" - # Ensure we're in a virtual environment - if not os.environ.get("VIRTUAL_ENV") and not args.no_venv_check: - print("⚠️ Not running in a virtual environment. Activate one before updating dependencies.") - print("💡 Run: uv venv && source .venv/bin/activate (or .venv\\Scripts\\activate on Windows)") - if not args.force: + if not check_uv_installed(): + return False + + print("Updating dependencies...") + + # Base dependencies + try: + cmd = ["uv", "pip", "install", "--upgrade", "-r", "requirements.txt"] + if lock: + cmd.append("--lock") + + subprocess.run(cmd, check=True) + print("✓ Base dependencies updated successfully.") + except subprocess.SubprocessError: + print("Failed to update base dependencies.") + return False + + # Advanced dependencies if requested + if advanced: + try: + if os.path.exists("requirements_advanced.txt"): + cmd = ["uv", "pip", "install", "--upgrade", "-r", "requirements_advanced.txt"] + if lock: + cmd.append("--lock") + + subprocess.run(cmd, check=True) + print("✓ Advanced dependencies updated successfully.") + else: + print("Advanced requirements file not found.") + except subprocess.SubprocessError: + print("Failed to update advanced dependencies.") return False - print("📦 Updating dependencies using UV...") - - # Update base requirements - subprocess.run(["uv", "pip", "install", "--upgrade", "-r", "requirements.txt"], check=True) - print("✅ Base dependencies updated") - - # Update advanced requirements if requested - if args.advanced: - subprocess.run(["uv", "pip", "install", "--upgrade", "-r", "requirements_advanced.txt"], check=True) - print("✅ Advanced dependencies updated") - - # Update PyTorch with CUDA if requested - if args.cuda: - torch_version = "2.4.0" if args.latest else "2.3.0" - cuda_version = "12.1" if args.latest else "11.8" + # CUDA dependencies if requested + if cuda: + print("Updating CUDA dependencies...") + extras = ["gpu"] + if advanced: + extras.append("advanced") - print(f"🔄 Installing PyTorch {torch_version} with CUDA {cuda_version}...") - subprocess.run([ - "uv", "pip", "install", "--upgrade", - f"torch=={torch_version}", - f"torchvision>={torch_version}", - f"torchaudio>={torch_version}", - "--index-url", f"https://download.pytorch.org/whl/cu{cuda_version.replace('.', '')}" - ], check=True) - print("✅ PyTorch updated with CUDA support") - - # Generate lockfile for reproducibility - if args.lock: - print("🔒 Generating lockfile...") - subprocess.run(["uv", "pip", "freeze", ">", "UV-req-lock.txt"], shell=True, check=True) - print("✅ Lockfile generated: UV-req-lock.txt") + try: + cmd = ["uv", "pip", "install", "--upgrade", f".{{{','.join(extras)}}}"] + if lock: + cmd.append("--lock") + + subprocess.run(cmd, check=True) + print("✓ CUDA dependencies updated successfully.") + except subprocess.SubprocessError: + print("Failed to update CUDA dependencies.") + return False return True def main(): - parser = argparse.ArgumentParser(description="Update ComfyUI dependencies using UV") + parser = argparse.ArgumentParser(description="Update ComfyUI dependencies with UV") parser.add_argument("--advanced", action="store_true", help="Update advanced dependencies") - parser.add_argument("--cuda", action="store_true", help="Update PyTorch with CUDA support") - parser.add_argument("--latest", action="store_true", help="Use latest versions (may be less stable)") - parser.add_argument("--lock", action="store_true", help="Generate lockfile after update") - parser.add_argument("--force", action="store_true", help="Force update even if not in venv") - parser.add_argument("--no-venv-check", action="store_true", help="Skip virtual environment check") + parser.add_argument("--cuda", action="store_true", help="Update CUDA dependencies") + parser.add_argument("--lock", action="store_true", help="Generate lockfile during update") args = parser.parse_args() - print("ComfyUI Dependency Update") - print("========================") - - if not check_uv_installed(): - sys.exit(1) - - if update_dependencies(args): - print("\n✨ Dependencies updated successfully!") - print("\n💡 To start ComfyUI, run: python main.py") + if update_dependencies(args.advanced, args.cuda, args.lock): + print("\n✓ All dependencies updated successfully!") else: - print("\n❌ Dependency update failed or was skipped") + print("\n❌ Failed to update some dependencies.") sys.exit(1) if __name__ == "__main__":