Merge pull request #18 from tweak-wtf/enhancement/dynamic-torch-pypi-index

Dynamic torch pypi index urls

Author: tweak-wtf

client/ayon_comfyui/hooks/pre_launch.py

@@ -82,6 +82,49 @@ def pre_process(self):
             cache_tmpl = self.addon_settings["caching"]["cache_dir_template"]
             self.cache_dir = StringTemplate(cache_tmpl).format_strict(self.tmpl_data)
 
+        # get installed CUDA version and build correct pypi index url
+        try:
+            smi_version_details = subprocess.check_output(
+                ["nvidia-smi", "--version"], text=True
+            ).strip()
+        except subprocess.CalledProcessError as e:
+            log.error(f"Failed to execute `nvidia-smi`: {e} Please ensure NVIDIA drivers are installed.")
+        
+        cuda_version = None
+        for line in smi_version_details.splitlines():
+            if "CUDA Version" in line:
+                parts = line.split(":")
+                cuda_version = parts[1].strip()
+                break
+        if not cuda_version:
+            log.error("Could not determine CUDA version from `nvidia-smi` output.")
+            raise RuntimeError("CUDA version could not be determined.")
+
+        pypi_url_map = {
+            "11.8": {
+                "stable": "https://download.pytorch.org/whl/cu118",
+                "nightly": None,
+            },
+            "12.6": {
+                "stable": "https://download.pytorch.org/whl/cu126",
+                "nightly": "https://download.pytorch.org/whl/nightly/cu126",
+            },
+            "12.8": {
+                "stable": "https://download.pytorch.org/whl/cu128",
+                "nightly": "https://download.pytorch.org/whl/nightly/cu128",
+            },
+            "12.9": {
+                "stable": None,
+                "nightly": "https://download.pytorch.org/whl/nightly/cu129",
+            },
+        }
+        if bool(self.addon_settings["venv"]["use_torch_nightly"]):
+            self.pypi_url = pypi_url_map[cuda_version]["nightly"]
+        else:
+            self.pypi_url = pypi_url_map[cuda_version]["stable"]
+
+        self.py_version = self.addon_settings["venv"]["python_version"]
+
     def clone_repositories(self):
         def git_clone(url: str, dest: Path, tag: str = "") -> git.Repo:
             if not dest.exists():
@@ -182,6 +225,12 @@ def run_server(self):
         if self.extra_flags:
             launch_args.append("-extraFlags")
             launch_args.append(",".join(self.extra_flags))
+        if self.pypi_url:
+            launch_args.append("-pypiUrl")
+            launch_args.append(self.pypi_url)
+        if self.py_version:
+            launch_args.append("-pythonVersion")
+            launch_args.append(self.py_version)
 
         _cmd.extend(launch_args)
         cmd = " ".join([str(arg) for arg in _cmd])

client/ayon_comfyui/tools/install_and_run_server_venv.ps1

@@ -1,6 +1,8 @@
 # assumes to be in comfyui directory
 param(
     [string]$cacheDir = "",
+    [string]$pypiUrl = "",
+    [string]$pythonVersion = "",
     [string[]]$plugins = @(),
     [string[]]$extraFlags = @(),
     [string[]]$extraDependencies = @()
@@ -18,16 +20,18 @@ if ($cacheDir) {
 }
 
 # create local venv
-uv venv --allow-existing --python 3.12
+uv venv --allow-existing --python $pythonVersion
 if (-not $?){
     Write-Output "Failed to create venv"
     exit 1
 }
 .venv\Scripts\activate
 
 # Install requirements
+Write-Output "Installing PyTorch with CUDA support"
+uv pip install --pre torch torchvision torchaudio --index-url $pypiUrl
 Write-Output "Installing ComfyUI requirements"
-uv pip install -r requirements.txt --extra-index-url https://download.pytorch.org/whl/cu126
+uv pip install -r requirements.txt
 
 # install plugins dependencies
 foreach ($plugin in $plugins) {

server/settings.py

@@ -18,6 +18,19 @@ def __init__(self, **data):
         self.name = Path(self.url).stem
 
 
+class VirtualEnvSettings(BaseSettingsModel):
+    python_version: str = SettingsField(
+        default="3.12",
+        title="Python Version",
+        description="Python version to use for the virtual environment.",
+    )
+    use_torch_nightly: bool = SettingsField(
+        default=True,
+        title="Use PyTorch Nightly",
+        description="Use the nightly version of PyTorch.",
+    )
+
+
 class CustomNodeSettings(RepositorySettings):
     extra_dependencies: list[str] = SettingsField(
         default_factory=list,
@@ -76,6 +89,11 @@ class AddonSettings(BaseSettingsModel):
         title="Extra Flags",
         description="Extra argument flags to pass when launching the ComfyUI server.",
     )
+    venv: VirtualEnvSettings = SettingsField(
+        default_factory=VirtualEnvSettings,
+        title="Virtual Environment Settings",
+        description="Virtual Environment Settings for the ComfyUI server.",
+    )
     repositories: ComfyUIRepositorySettings = SettingsField(
         default_factory=ComfyUIRepositorySettings,
         title="Repository Settings",