Update to torch 2.0.0 and the pre-release xformers library

requirements.txt

@@ -1,23 +1,23 @@
-torch==1.13.1
-torchvision==0.14.1
-torchaudio==0.13.1
-torchdiffeq==0.2.3
-torchsde==0.2.5
-einops==0.6.0
-open-clip-torch==2.16.0
+torch>=1.13.0
+torchvision
+torchaudio
+torchdiffeq>=0.2.3
+torchsde>=0.2.5
+einops>=0.6.0
+open-clip-torch>=2.16.0
 transformers>=4.25.1
 safetensors>=0.3.0
-pytorch-lightning==2.0.0
-aiohttp==3.8.4
-accelerate==0.18.0
-pyyaml==6.0
-scikit-image==0.20.0
-jsonmerge==1.9.0
-clean-fid==0.1.35
-clip==0.2.0
-resize-right==0.0.2
-opencv-python==4.7.0.72
-albumentations==1.3.0
+pytorch-lightning>=2.0.0
+aiohttp>=3.8.4
+accelerate>=0.18.0
+pyyaml>=6.0
+scikit-image>=0.20.0
+jsonmerge>=1.9.0
+clean-fid>=0.1.35
+clip>=0.2.0
+resize-right>=0.0.2
+opencv-python>=4.7.0.72
+albumentations>=1.3.0
 Pillow
 scipy
 tqdm

setup.py

@@ -3,6 +3,7 @@
 import os.path
 import platform
 import subprocess
+import sys
 
 from pip._internal.index.collector import LinkCollector
 from pip._internal.index.package_finder import PackageFinder
@@ -13,9 +14,40 @@ from pip._internal.req import InstallRequirement
 from pip._vendor.packaging.requirements import Requirement
 from setuptools import setup, find_packages
 
+"""
+The name of the package.
+"""
 package_name = "comfyui"
+
+"""
+The current version.
+"""
 version = '0.0.1'
 
+"""
+The package index to the torch built with AMD ROCm.
+"""
+amd_torch_index = "https://download.pytorch.org/whl/rocm5.4.2"
+
+"""
+The package index to torch built with CUDA.
+Observe the CUDA version is in this URL.
+"""
+nvidia_torch_index = "https://download.pytorch.org/whl/cu118"
+
+"""
+The package index to torch built against CPU features.
+This includes macOS MPS support.
+"""
+cpu_torch_index_nightlies = "https://download.pytorch.org/whl/nightly/cpu"
+
+"""
+The xformers dependency and version string.
+This should be updated whenever another pre-release of xformers is supported. The current build was retrieved from
+https://pypi.org/project/xformers/0.0.17rc482/#history.
+"""
+xformers_dep = "xformers==0.0.17rc482"
+
 
 def _is_nvidia() -> bool:
     system = platform.system().lower()
@@ -60,29 +92,38 @@ def _is_amd() -> bool:
     return False
 
 
-_amd_torch_index = "https://download.pytorch.org/whl/rocm5.4.2"
-_nvidia_torch_index = "https://download.pytorch.org/whl/cu117"
-_alternative_indices = [_amd_torch_index, _nvidia_torch_index]
-
-
 def dependencies() -> [str]:
     _dependencies = open(os.path.join(os.path.dirname(__file__), "requirements.txt")).readlines()
-
+    _alternative_indices = [amd_torch_index, nvidia_torch_index, cpu_torch_index_nightlies]
     session = PipSession()
 
     index_urls = ['https://pypi.org/simple']
     # prefer nvidia over AMD because AM5/iGPU systems will have a valid ROCm device
     if _is_nvidia():
-        index_urls += [_nvidia_torch_index]
-        _dependencies += ["xformers==0.0.16"]
+        index_urls += [nvidia_torch_index]
+        _dependencies += [xformers_dep]
     elif _is_amd():
-        index_urls += [_amd_torch_index]
+        index_urls += [amd_torch_index]
+    else:
+        index_urls += [cpu_torch_index_nightlies]
 
     if len(index_urls) == 1:
         return _dependencies
-    finder = PackageFinder.create(LinkCollector(session, SearchScope([], index_urls, no_index=False)),
-                                  SelectionPreferences(allow_yanked=False, prefer_binary=False))
 
+    try:
+        # pip 23
+        finder = PackageFinder.create(LinkCollector(session, SearchScope([], index_urls, no_index=False)),
+                                      SelectionPreferences(allow_yanked=False, prefer_binary=False,
+                                                           allow_all_prereleases=True))
+    except:
+        try:
+            # pip 22
+            finder = PackageFinder.create(LinkCollector(session, SearchScope([], index_urls)),
+                                          SelectionPreferences(allow_yanked=False, prefer_binary=False,
+                                                               allow_all_prereleases=True)
+                                          , use_deprecated_html5lib=False)
+        except:
+            raise Exception("upgrade pip with\npip install -U pip")
     for i, package in enumerate(_dependencies[:]):
         requirement = InstallRequirement(Requirement(package), comes_from=f"{package_name}=={version}")
         candidate = finder.find_best_candidate(requirement.name, requirement.specifier)