refactor: enhance dependency update script with improved UV installation checks and error handling

update_dependencies_uv.py

@@ -1,103 +1,110 @@
 #!/usr/bin/env python3
 """
-ComfyUI Dependency Update Script using UV
+Update dependencies for ComfyUI using UV
-Created: May 2, 2025
 """
 
-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 check_uv_installed():
-    if not os.environ.get("VIRTUAL_ENV") and not args.no_venv_check:
+        return False
-        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)")
+    print("Updating dependencies...")
-        if not args.force:
+    
+    # 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...")
+    # CUDA dependencies if requested
-    
+    if cuda:
-    # Update base requirements
+        print("Updating CUDA dependencies...")
-    subprocess.run(["uv", "pip", "install", "--upgrade", "-r", "requirements.txt"], check=True)
+        extras = ["gpu"]
-    print("✅ Base dependencies updated")
+        if advanced:
-    
+            extras.append("advanced")
-    # 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"
         
-        print(f"🔄 Installing PyTorch {torch_version} with CUDA {cuda_version}...")
+        try:
-        subprocess.run([
+            cmd = ["uv", "pip", "install", "--upgrade", f".{{{','.join(extras)}}}"]
-            "uv", "pip", "install", "--upgrade", 
+            if lock:
-            f"torch=={torch_version}", 
+                cmd.append("--lock")
-            f"torchvision>={torch_version}", 
+            
-            f"torchaudio>={torch_version}", 
+            subprocess.run(cmd, check=True)
-            "--index-url", f"https://download.pytorch.org/whl/cu{cuda_version.replace('.', '')}"
+            print("✓ CUDA dependencies updated successfully.")
-        ], check=True)
+        except subprocess.SubprocessError:
-        print("✅ PyTorch updated with CUDA support")
+            print("Failed to update CUDA dependencies.")
-    
+            return False
-    # 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")
     
     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("--cuda", action="store_true", help="Update CUDA dependencies")
-    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 during update")
-    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")
     
     args = parser.parse_args()
     
-    print("ComfyUI Dependency Update")
+    if update_dependencies(args.advanced, args.cuda, args.lock):
-    print("========================")
+        print("\n✓ All dependencies updated successfully!")
-    
-    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")
     else:
-        print("\n❌ Dependency update failed or was skipped")
+        print("\n❌ Failed to update some dependencies.")
         sys.exit(1)
 
 if __name__ == "__main__":