removed str from Union[GPU_OPTION, str, None] typing in a bunch o…

…f places
Comfy-Org · Aug 24, 2024 · 2a1762c · 2a1762c
1 parent 08c553f
commit 2a1762c
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 9 deletions.
diff --git a/comfy_cli/standalone.py b/comfy_cli/standalone.py
@@ -7,7 +7,7 @@
 
 import requests
 
-from comfy_cli.constants import OS, PROC
+from comfy_cli.constants import GPU_OPTION, OS, PROC
 from comfy_cli.typing import PathLike
 from comfy_cli.utils import download_progress, get_os, get_proc
 from comfy_cli.uv import DependencyCompiler
@@ -92,6 +92,8 @@ def FromTarball(fpath: PathLike, name: PathLike = "python"):
 
         old_rpath = fpath.parent / old_name
         rpath = fpath.parent / name
+        # clean the expanded destination
+        shutil.rmtree(rpath, ignore_errors=True)
         shutil.move(old_rpath, rpath)
         return StandalonePython(rpath=rpath)
 
@@ -138,25 +140,25 @@ def run_comfy_cli(self, *args: list[str]):
     def install_comfy(self, *args: list[str], gpu_arg: str = "--nvidia"):
         self.run_comfy_cli("--here", "--skip-prompt", "install", "--fast-deps", gpu_arg, *args)
 
-    def compile_comfy_deps(self, comfyDir: PathLike, gpu: str, outDir: Optional[PathLike] = None):
+    def compile_comfy_deps(self, comfyDir: PathLike, gpu: GPU_OPTION, outDir: Optional[PathLike] = None):
         outDir = self.rpath if outDir is None else outDir
 
         self.dep_comp = DependencyCompiler(cwd=comfyDir, executable=self.executable, gpu=gpu, outDir=outDir)
         self.dep_comp.compile_comfy_deps()
 
-    def install_comfy_deps(self, comfyDir: PathLike, gpu: str, outDir: Optional[PathLike] = None):
+    def install_comfy_deps(self, comfyDir: PathLike, gpu: GPU_OPTION, outDir: Optional[PathLike] = None):
         outDir = self.rpath if outDir is None else outDir
 
         self.dep_comp = DependencyCompiler(cwd=comfyDir, executable=self.executable, gpu=gpu, outDir=outDir)
         self.dep_comp.install_core_plus_ext()
 
-    def precache_comfy_deps(self, comfyDir: PathLike, gpu: str, outDir: Optional[PathLike] = None):
+    def precache_comfy_deps(self, comfyDir: PathLike, gpu: GPU_OPTION, outDir: Optional[PathLike] = None):
         outDir = self.rpath if outDir is None else outDir
 
         self.dep_comp = DependencyCompiler(cwd=comfyDir, executable=self.executable, gpu=gpu, outDir=outDir)
         self.dep_comp.precache_comfy_deps()
 
-    def wheel_comfy_deps(self, comfyDir: PathLike, gpu: str, outDir: Optional[PathLike] = None):
+    def wheel_comfy_deps(self, comfyDir: PathLike, gpu: GPU_OPTION, outDir: Optional[PathLike] = None):
         outDir = self.rpath if outDir is None else outDir
 
         self.dep_comp = DependencyCompiler(cwd=comfyDir, executable=self.executable, gpu=gpu, outDir=outDir)

diff --git a/comfy_cli/uv.py b/comfy_cli/uv.py
@@ -265,7 +265,7 @@ def Wheel(
         return _check_call(cmd, cwd)
 
     @staticmethod
-    def Resolve_Gpu(gpu: Union[GPU_OPTION, str, None]):
+    def Resolve_Gpu(gpu: Union[GPU_OPTION, None]):
         if gpu is None:
             try:
                 tver = metadata.version("torch")
@@ -277,16 +277,14 @@ def Resolve_Gpu(gpu: Union[GPU_OPTION, str, None]):
                     return None
             except metadata.PackageNotFoundError:
                 return None
-        elif isinstance(gpu, str):
-            return GPU_OPTION[gpu.upper()]
         else:
             return gpu
 
     def __init__(
         self,
         cwd: PathLike = ".",
         executable: PathLike = sys.executable,
-        gpu: Union[GPU_OPTION, str, None] = None,
+        gpu: Union[GPU_OPTION, None] = None,
         outDir: PathLike = ".",
         outName: str = "requirements.compiled",
         reqFilesCore: Optional[list[PathLike]] = None,