Create a setup.py and automatically select the correct pytorch binaries for the current platform and supported devices

- setup.py now works
 - Makes installation work a variety of ways, including making other packages dependent on this one for e.g. plugins
 - Fixes missing __init__.py issues
 - Fixes imports
 - Compatible with your existing scripts that rely on requirements.txt
 - Fixes error in comfy/ldm/models/diffusion/ddim.py
 - Fixes missing packages for other diffusers code in this repo

.gitignore

@@ -1,5 +1,3 @@
-__pycache__/
-*.py[cod]
 output/
 input/
 !input/example.png
@@ -14,3 +12,158 @@ venv/
 web/extensions/*
 !web/extensions/logging.js.example
 !web/extensions/core/
+
+# The GitHub python gitignore
+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.py[cod]
+*$py.class
+
+# C extensions
+*.so
+
+# Distribution / packaging
+.Python
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+wheels/
+share/python-wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+MANIFEST
+
+# PyInstaller
+#  Usually these files are written by a python script from a template
+#  before PyInstaller builds the exe, so as to inject date/other infos into it.
+*.manifest
+*.spec
+
+# Installer logs
+pip-log.txt
+pip-delete-this-directory.txt
+
+# Unit test / coverage reports
+htmlcov/
+.tox/
+.nox/
+.coverage
+.coverage.*
+.cache
+nosetests.xml
+coverage.xml
+*.cover
+*.py,cover
+.hypothesis/
+.pytest_cache/
+cover/
+
+# Translations
+*.mo
+*.pot
+
+# Django stuff:
+*.log
+local_settings.py
+db.sqlite3
+db.sqlite3-journal
+
+# Flask stuff:
+instance/
+.webassets-cache
+
+# Scrapy stuff:
+.scrapy
+
+# Sphinx documentation
+docs/_build/
+
+# PyBuilder
+.pybuilder/
+target/
+
+# Jupyter Notebook
+.ipynb_checkpoints
+
+# IPython
+profile_default/
+ipython_config.py
+
+# pyenv
+#   For a library or package, you might want to ignore these files since the code is
+#   intended to run in multiple environments; otherwise, check them in:
+# .python-version
+
+# pipenv
+#   According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
+#   However, in case of collaboration, if having platform-specific dependencies or dependencies
+#   having no cross-platform support, pipenv may install dependencies that don't work, or not
+#   install all needed dependencies.
+#Pipfile.lock
+
+# poetry
+#   Similar to Pipfile.lock, it is generally recommended to include poetry.lock in version control.
+#   This is especially recommended for binary packages to ensure reproducibility, and is more
+#   commonly ignored for libraries.
+#   https://python-poetry.org/docs/basic-usage/#commit-your-poetrylock-file-to-version-control
+#poetry.lock
+
+# pdm
+#   Similar to Pipfile.lock, it is generally recommended to include pdm.lock in version control.
+#pdm.lock
+#   pdm stores project-wide configurations in .pdm.toml, but it is recommended to not include it
+#   in version control.
+#   https://pdm.fming.dev/#use-with-ide
+.pdm.toml
+
+# PEP 582; used by e.g. github.com/David-OConnor/pyflow and github.com/pdm-project/pdm
+__pypackages__/
+
+# Celery stuff
+celerybeat-schedule
+celerybeat.pid
+
+# SageMath parsed files
+*.sage.py
+
+# Environments
+.env
+.venv
+env/
+venv/
+ENV/
+env.bak/
+venv.bak/
+
+# Spyder project settings
+.spyderproject
+.spyproject
+
+# Rope project settings
+.ropeproject
+
+# mkdocs documentation
+/site
+
+# mypy
+.mypy_cache/
+.dmypy.json
+dmypy.json
+
+# Pyre type checker
+.pyre/
+
+# pytype static type analyzer
+.pytype/
+
+# Cython debug symbols
+cython_debug/

README.md

@@ -88,19 +88,13 @@ Put your SD checkpoints (the huge ckpt/safetensors files) in: models/checkpoints
 
 Put your VAE in: models/vae
 
-### AMD GPUs (Linux only)
-AMD users can install rocm and pytorch with pip if you don't have it already installed, this is the command to install the stable version:
+Then, run the following command to install comfyui into your current environment:
 
-```pip install torch torchvision torchaudio --extra-index-url https://download.pytorch.org/whl/rocm5.4.2```
+```
+pip install -e .
+```
 
-This is the command to install the nightly with ROCm 5.6 that supports the 7000 series and might have some performance improvements:
-```pip install --pre torch torchvision torchaudio --index-url https://download.pytorch.org/whl/nightly/rocm5.6```
-
-### NVIDIA
-
-Nvidia users should install torch and xformers using this command:
-
-```pip install torch torchvision torchaudio --extra-index-url https://download.pytorch.org/whl/cu118 xformers```
+This will correctly select the version of pytorch that matches the GPU on your machine (NVIDIA or CPU on Windows, NVIDIA AMD or CPU on Linux).
 
 #### Troubleshooting
 
@@ -110,14 +104,6 @@ If you get the "Torch not compiled with CUDA enabled" error, uninstall torch wit
 
 And install it again with the command above.
 
-### Dependencies
-
-Install the dependencies by opening your terminal inside the ComfyUI folder and:
-
-```pip install -r requirements.txt```
-
-After this you should have everything installed and can proceed to running ComfyUI.
-
 ### Others:
 
 #### [Intel Arc](https://github.com/comfyanonymous/ComfyUI/discussions/476)

comfy/extra_samplers/uni_pc.py

@@ -688,7 +688,7 @@ class UniPC:
                 x_t = x_t_ - expand_dims(alpha_t * B_h, dims) * (corr_res + rhos_c[-1] * D1_t)
         else:
             x_t_ = (
-                expand_dims(torch.exp(log_alpha_t - log_alpha_prev_0), dimss) * x
+                expand_dims(torch.exp(log_alpha_t - log_alpha_prev_0), dims) * x
                 - expand_dims(sigma_t * h_phi_1, dims) * model_prev_0
             )
             if x_t is None:

comfy/ldm/models/diffusion/plms.py

@@ -5,8 +5,8 @@ import numpy as np
 from tqdm import tqdm
 from functools import partial
 
-from ldm.modules.diffusionmodules.util import make_ddim_sampling_parameters, make_ddim_timesteps, noise_like
-from ldm.models.diffusion.sampling_util import norm_thresholding
+from comfy.ldm.modules.diffusionmodules.util import make_ddim_sampling_parameters, make_ddim_timesteps, noise_like
+from comfy.ldm.models.diffusion.sampling_util import norm_thresholding
 
 
 class PLMSSampler(object):

requirements.txt

@@ -1,11 +1,23 @@
-torch
-torchsde
-einops
+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
 transformers>=4.25.1
 safetensors>=0.3.0
-aiohttp
-accelerate
-pyyaml
+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
+tqdm

setup.py

@@ -0,0 +1,103 @@
+#!/usr/bin/env python3
+# this script does a little housekeeping for your platform
+import os.path
+import platform
+import subprocess
+
+from pip._internal.index.collector import LinkCollector
+from pip._internal.index.package_finder import PackageFinder
+from pip._internal.models.search_scope import SearchScope
+from pip._internal.models.selection_prefs import SelectionPreferences
+from pip._internal.network.session import PipSession
+from pip._internal.req import InstallRequirement
+from pip._vendor.packaging.requirements import Requirement
+from setuptools import setup, find_packages
+
+package_name = "comfyui"
+version = '0.0.1'
+
+
+def _is_nvidia() -> bool:
+    system = platform.system().lower()
+    nvidia_smi_paths = []
+
+    if system == "windows":
+        nvidia_smi_paths.append(os.path.join(os.environ.get("SystemRoot", ""), "System32", "nvidia-smi.exe"))
+    elif system == "linux":
+        nvidia_smi_paths.extend(["/usr/bin/nvidia-smi", "/opt/nvidia/bin/nvidia-smi"])
+
+    for nvidia_smi_path in nvidia_smi_paths:
+        try:
+            output = subprocess.check_output([nvidia_smi_path, "-L"]).decode("utf-8")
+
+            if "GPU" in output:
+                return True
+        except:
+            pass
+
+    return False
+
+
+def _is_amd() -> bool:
+    system = platform.system().lower()
+    rocminfo_paths = []
+
+    # todo: torch windows doesn't support amd
+    if system == "windows":
+        rocminfo_paths.append(os.path.join(os.environ.get("ProgramFiles", ""), "AMD", "ROCm", "bin", "rocminfo.exe"))
+    elif system == "linux":
+        rocminfo_paths.extend(["/opt/rocm/bin/rocminfo", "/usr/bin/rocminfo"])
+
+    for rocminfo_path in rocminfo_paths:
+        try:
+            output = subprocess.check_output([rocminfo_path]).decode("utf-8")
+
+            if "Device" in output:
+                return True
+        except:
+            pass
+
+    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()
+
+    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"]
+    elif _is_amd():
+        index_urls += [_amd_torch_index]
+
+    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))
+
+    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)
+        if any([url in candidate.best_candidate.link.url for url in _alternative_indices]):
+            _dependencies[i] = f"{requirement.name} @ {candidate.best_candidate.link.url}"
+    return _dependencies
+
+
+setup(
+    # "comfyui"
+    name=package_name,
+    description="",
+    author="",
+    version=version,
+    python_requires=">=3.9,<3.11",
+    packages=find_packages(include=['comfy', 'comfy_extras']),
+    install_requires=dependencies(),
+)