From 7970655b1dd1ba2e724a6106ecb054f3a8e51f71 Mon Sep 17 00:00:00 2001
From: telamonian <telamonian@users.noreply.github.com>
Date: Fri, 23 Aug 2024 03:16:52 -0400
Subject: [PATCH] added `install_*` methods to `StandalonePython`

---
 comfy_cli/standalone.py | 37 ++++++++++++++++++++++++++++++++++++-
 comfy_cli/uv.py         | 12 +++++-------
 pyproject.toml          |  1 +
 3 files changed, 42 insertions(+), 8 deletions(-)

diff --git a/comfy_cli/standalone.py b/comfy_cli/standalone.py
index eec8354..99221cd 100644
--- a/comfy_cli/standalone.py
+++ b/comfy_cli/standalone.py
@@ -1,14 +1,18 @@
 import os
 from pathlib import Path
 import shutil
+import subprocess
 import tarfile
 from typing import Optional
 
 import requests
 
 from comfy_cli.constants import OS, PROC
-from comfy_cli.utils import download_progress, get_os, get_proc
 from comfy_cli.typing import PathLike
+from comfy_cli.utils import download_progress, get_os, get_proc
+from comfy_cli.uv import DependencyCompiler
+
+_here = Path(__file__).expanduser().resolve().parent
 
 _platform_targets = {
     (OS.MACOS, PROC.ARM): "aarch64-apple-darwin",
@@ -85,3 +89,34 @@ def __init__(self, rpath: PathLike):
         self.name = self.rpath.name
         self.bin = self.rpath / "bin"
         self.executable = self.bin / "python"
+
+        # upgrade pip if needed, install uv
+        self.pip_install("-U", "pip", "uv")
+
+    def run_module(self, mod: str, *args: list[str]):
+        cmd: list[str] = [
+            str(self.executable),
+            "-m",
+            mod,
+            *args,
+        ]
+
+        subprocess.run(cmd, check=True)
+
+    def pip_install(self, *args: list[str]):
+        self.run_module("pip", "install", *args)
+
+    def uv_install(self, *args: list[str]):
+        self.run_module("uv", "pip", "install", *args)
+
+    def install_comfy_cli(self, dev: bool = False):
+        if dev:
+            self.uv_install(str(_here.parent))
+        else:
+            self.uv_install("comfy_cli")
+
+    def run_comfy_cli(self, *args: list[str]):
+        self.run_module("comfy_cli", *args)
+
+    def install_comfy(self, *args: list[str], gpu_arg: str = "--nvidia"):
+        self.run_comfy_cli("--here", "--skip-prompt", "install", "--fast-deps", gpu_arg, *args)
diff --git a/comfy_cli/uv.py b/comfy_cli/uv.py
index 4d7f505..904af10 100644
--- a/comfy_cli/uv.py
+++ b/comfy_cli/uv.py
@@ -12,8 +12,8 @@
 from comfy_cli.typing import PathLike
 
 
-def _run(cmd: list[str], cwd: PathLike) -> subprocess.CompletedProcess[Any]:
-    return subprocess.run(cmd, cwd=cwd, capture_output=True, text=True, check=True)
+def _run(cmd: list[str], cwd: PathLike, check: bool = True) -> subprocess.CompletedProcess[Any]:
+    return subprocess.run(cmd, cwd=cwd, capture_output=True, text=True, check=check)
 
 
 def _check_call(cmd: list[str], cwd: Optional[PathLike] = None):
@@ -45,7 +45,7 @@ def parse_uv_compile_error(err: str) -> tuple[str, list[str]]:
 
 
 class DependencyCompiler:
-    rocmPytorchUrl = "https://download.pytorch.org/whl/rocm6.0"
+    rocmPytorchUrl = "https://download.pytorch.org/whl/rocm6.1"
     nvidiaPytorchUrl = "https://download.pytorch.org/whl/cu121"
 
     overrideGpu = dedent(
@@ -77,10 +77,8 @@ def Find_Req_Files(*ders: PathLike) -> list[Path]:
     @staticmethod
     def Install_Build_Deps(executable: PathLike = sys.executable):
         """Use pip to install bare minimum requirements for uv to do its thing"""
-        if shutil.which("uv") is None:
-            cmd = [str(executable), "-m", "pip", "install", "--upgrade", "pip", "uv"]
-
-            _check_call(cmd=cmd)
+        cmd = [str(executable), "-m", "pip", "install", "--upgrade", "pip", "uv"]
+        _check_call(cmd=cmd)
 
     @staticmethod
     def Compile(
diff --git a/pyproject.toml b/pyproject.toml
index 78f20ca..f8a21a7 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -40,6 +40,7 @@ dependencies = [
   "tqdm",
   "typer>=0.9.0",
   "typing-extensions>=4.7.0",
+  "uv",
   "websocket-client",
 ]