93 Commits
Author SHA1 Message Date
Limitex 3cabac3ea7 refactor: scheduler type map 2025-11-12 22:37:37 +09:00
Limitex f7292b3c96 feat: add prompt to conditioning 2025-11-12 22:04:55 +09:00
Limitex 99f2c32d08 fix: scheduler dto 2025-11-12 22:03:52 +09:00
Limitex bf99ccfc02 fix(node): image size limit 2025-11-12 21:53:48 +09:00
Limitex 0fced1b678 refactor: transferring cache to infra 2025-11-12 21:53:48 +09:00
Limitex b4d197809c feat(domain): add validation and encapsulation to domain models 2025-11-12 21:53:48 +09:00
Limitex c4d385be60 refactor: validation 2025-11-09 09:21:30 +09:00
Limitex f6fb30d5aa refactor: merge Handler and Service layers into UseCase 2025-11-09 09:16:35 +09:00
Limitex 734207e52c refactor(infra): replace custom VAE loader with Diffusers native API 2025-11-09 08:27:18 +09:00
Limitex 73e4c8dbda fix: cache path for diffusers repository cache 2025-11-09 07:51:13 +09:00
Limitex b29acb6534 feat: get cheche dir paths 2025-11-09 07:51:13 +09:00
Limitex f7d6e25766 feat: get project root config 2025-11-09 07:08:05 +09:00
Limitex d31484ff5a feat: cache dir env 2025-11-09 07:08:05 +09:00
Limitex 156bc90fc0 faet: dotenv 2025-11-09 06:57:32 +09:00
Limitex fef82b2c85 vscode: add cspell 2025-11-09 05:58:32 +09:00
Limitex bb0c1ae3e1 add: python-dotenv 2025-11-09 05:29:38 +09:00
Limitex eeacae5b7c feat: add cache 2025-11-09 04:54:54 +09:00
Limitex 69608c7eab fix: taskfile structure 2025-11-09 01:28:26 +09:00
Limitex 975cfa7d0b chore: add tree to taskfile 2025-11-09 01:10:37 +09:00
Limitex da895fb3ec chore: add lockup to taskfile 2025-11-09 00:45:28 +09:00
Limitex dd5ccc0395 chore: update taskfile 2025-11-09 00:42:08 +09:00
Limitex 86e365da23 Merge branch 'develop/docker' into develop/replace 2025-11-08 23:27:15 +09:00
Limitex 25d36656a7 fix: taskfile 2025-11-08 23:17:46 +09:00
Limitex 714251ca99 chore: sync requirements 2025-11-08 23:12:01 +09:00
Limitex 1686e9fa1e chore: docker for developing 2025-11-08 22:47:55 +09:00
Limitex 40a9ec8227 chore: use cuda 12.4 2025-11-08 21:57:41 +09:00
Limitex ca8d64fd84 fix: clip domain 2025-11-05 22:42:56 +09:00
Limitex 8399839699 chore: move files 2025-11-05 22:36:06 +09:00
Limitex 26a12aab12 chore: connect nodes to sampler 2025-11-05 22:31:15 +09:00
Limitex e45ba2ece7 feat: to domain dto 2025-11-05 22:10:53 +09:00
Limitex 4be7637b51 feat: scheduler node 2025-11-05 22:06:01 +09:00
Limitex 2f4a4d39fb vscode: add cspell 2025-11-04 08:59:11 +09:00
Limitex 28794f216d vscode: add cspell 2025-11-03 21:02:06 +09:00
Limitex bd32d3e7f2 feat: sampler node 2025-11-03 21:02:05 +09:00
Limitex 98b85781f3 refactor: architecture 2025-11-03 06:02:16 +09:00
Limitex 955fa8a7fa feat: clip text encode 2025-11-03 05:44:37 +09:00
Limitex 9c6cd87af9 chore: export clip from pipeline 2025-11-02 23:46:01 +09:00
Limitex f0f22cab08 feat: clip model 2025-11-02 23:05:21 +09:00
Limitex c87630816b feat: device load 2025-11-02 22:27:24 +09:00
Limitex c96edd2c53 chore(mypy): add rules 2025-11-02 22:08:14 +09:00
Limitex b4df7a2128 refactor: clarifying responsibility boundaries for folder_paths 2025-11-02 20:50:59 +09:00
Limitex e463490409 chore: add gitignore 2025-11-01 22:06:17 +09:00
Limitex c2bc6996ad chore: add cache to tmp dir 2025-11-01 22:03:49 +09:00
Limitex 19ba177059 chore: add dockerignore 2025-11-01 22:01:37 +09:00
Limitex 66270d0a1b feat: create dockerignore 2025-11-01 21:59:28 +09:00
Limitex 99947be37f lint: add mypy ignore 2025-11-01 21:44:42 +09:00
Limitex 9234093723 update: poetry lock 2025-11-01 21:39:48 +09:00
Limitex 1c2d8baec3 chore(mypy): add main module import missing rule 2025-11-01 20:49:19 +09:00
Limitex 681497608a remove: echo from taskfile 2025-11-01 20:41:31 +09:00
Limitex e5149c6f32 feat(docker): project dependence 2025-11-01 20:35:29 +09:00
Limitex 9a96665393 refactor(docker): optimize dockerfile 2025-11-01 13:08:31 +09:00
Limitex e6e3ae8f70 feat(docker): develop environment 2025-11-01 12:33:07 +09:00
Limitex df3e419071 feat(docker): comfyui docker container 2025-11-01 11:44:30 +09:00
Limitex d44e31b3f8 vscode: cspell settings 2025-11-01 00:39:47 +09:00
Limitex b7f837f96a fix: format check task 2025-11-01 00:39:47 +09:00
Limitex 1f9028a6f0 style: rename 2025-11-01 00:39:46 +09:00
Limitex 7a94a706d8 chore: add warning 2025-10-31 23:42:59 +09:00
Limitex 6de3b9ab6e style: divide into files 2025-10-31 23:33:45 +09:00
Limitex 68173980b5 fix: format 2025-10-31 23:29:07 +09:00
Limitex 9ba658fc46 update: linter and formatter settings 2025-10-31 23:27:24 +09:00
Limitex d5f0c805e2 add: commitizen 2025-10-31 22:39:47 +09:00
Limitex dc0d3d3572 add: pre-commit 2025-10-31 22:39:19 +09:00
Limitex 0ec12a27f8 add: Import-Linter 2025-10-31 22:38:20 +09:00
Limitex d72a2220f6 add: pytest 2025-10-31 22:08:54 +09:00
Limitex bb8621c7c5 update: ruff fsspec 2025-10-31 22:08:54 +09:00
Limitex 7e79e4c557 style: rename files 2025-10-31 21:12:55 +09:00
Limitex fd9d18f1ca remove: fast api 2025-10-31 21:03:18 +09:00
Limitex 1c96f63a4b feat: create vae loader 2025-10-31 20:53:51 +09:00
Limitex bcee329ca8 fix: taskfile 2025-10-31 20:53:33 +09:00
Limitex 7bb566d071 add: types-requests 2025-10-31 20:53:33 +09:00
Limitex 46b58feadb add: omegaconf 2025-10-31 20:53:33 +09:00
Limitex fc0d53e167 style: rename 2025-10-31 20:53:33 +09:00
Limitex d0c07557d8 fix: taskfile 2025-10-31 20:53:33 +09:00
Limitex c39d32f8cd feat: comfy dto 2025-10-31 20:53:33 +09:00
Limitex 59c94c6439 feat: create architecture with diffusers pipeline 2025-10-31 20:53:33 +09:00
Limitex 40a179f5d2 fix: src reference 2025-10-30 20:36:59 +09:00
Limitex bb00a7ee9d add: transformers 2025-10-30 08:01:48 +09:00
Limitex ead4abb8c9 add: torch 2025-10-30 08:01:48 +09:00
Limitex 98b4ea034c feat: add export requriements.txt to task file 2025-10-30 08:01:48 +09:00
Limitex c6df542fed add: fast api 2025-10-30 08:01:48 +09:00
Limitex be72f31045 add: diffusers 2025-10-30 08:01:48 +09:00
Limitex cb31f8526b add: dependency-injector 2025-10-30 08:01:48 +09:00
Limitex c79695c4b2 chore: init 2025-10-30 08:01:48 +09:00
Limitex 07d426f79d remove: legacy files 2025-10-29 21:41:49 +09:00
Limitex 4498af1394 Merge pull request #35 from Limitex/Limitex-patch-1
Update pyproject.toml
2025-02-20 21:39:50 +09:00
Limitex 87f089b6e1 Update pyproject.toml 2025-02-20 21:39:40 +09:00
Limitex 85b148d3fd Merge pull request #32 from doctorjei/main
Add lower-version limit to diffusers module to avoid conflicts with recent software (including ComfyUI-Manager)
2025-02-20 21:38:54 +09:00
Jeremiah Blanchard 7268b9f61c Update requirements.txt 2025-01-11 09:02:22 -05:00
Jeremiah Blanchard ec72e0fe52 Update pyproject.toml
ComfyUI-Manager depends on huggingface_hub>=0.20, which is incompatible with diffusers versions before 0.29.0. This prevents other packages from breaking ComfyUI-Manager when installing ComfyUI-Diffusers.
2025-01-11 09:01:59 -05:00
Limitex 8fbed904b4 Merge pull request #31 from Limitex/Limitex-patch-1
chore: update version
2025-01-09 14:30:40 +09:00
Limitex 5a2c333036 chore: update version 2025-01-09 14:30:26 +09:00
Limitex e83acdbb97 Merge pull request #30 from ComfyNodePRs/licence-update
Update PyProject Toml - License
2025-01-09 14:25:51 +09:00
snomiao 85b21e18d0 chore(licence-update): Update PyProject Toml - License 2025-01-09 02:00:19 +00:00
78 changed files with 5837 additions and 781 deletions
+264
View File
@@ -0,0 +1,264 @@
# Byte-compiled / optimized / DLL files
__pycache__/
*.py[codz]
*$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
# UV
# Similar to Pipfile.lock, it is generally recommended to include uv.lock in version control.
# This is especially recommended for binary packages to ensure reproducibility, and is more
# commonly ignored for libraries.
# uv.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
# poetry.toml
# pdm
# Similar to Pipfile.lock, it is generally recommended to include pdm.lock in version control.
# pdm recommends including project-wide configuration in pdm.toml, but excluding .pdm-python.
# https://pdm-project.org/en/latest/usage/project/#working-with-version-control
# pdm.lock
# pdm.toml
.pdm-python
.pdm-build/
# pixi
# Similar to Pipfile.lock, it is generally recommended to include pixi.lock in version control.
# pixi.lock
# Pixi creates a virtual environment in the .pixi directory, just like venv module creates one
# in the .venv directory. It is recommended not to include this directory in version control.
.pixi
# PEP 582; used by e.g. github.com/David-OConnor/pyflow and github.com/pdm-project/pdm
__pypackages__/
# Celery stuff
celerybeat-schedule
celerybeat.pid
# Redis
*.rdb
*.aof
*.pid
# RabbitMQ
mnesia/
rabbitmq/
rabbitmq-data/
# ActiveMQ
activemq-data/
# SageMath parsed files
*.sage.py
# Environments
.env
.envrc
.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/
# PyCharm
# JetBrains specific template is maintained in a separate JetBrains.gitignore that can
# be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore
# and can be added to the global gitignore or merged into this file. For a more nuclear
# option (not recommended) you can uncomment the following to ignore the entire idea folder.
# .idea/
# Abstra
# Abstra is an AI-powered process automation framework.
# Ignore directories containing user credentials, local state, and settings.
# Learn more at https://abstra.io/docs
.abstra/
# Visual Studio Code
# Visual Studio Code specific template is maintained in a separate VisualStudioCode.gitignore
# that can be found at https://github.com/github/gitignore/blob/main/Global/VisualStudioCode.gitignore
# and can be added to the global gitignore or merged into this file. However, if you prefer,
# you could uncomment the following to ignore the entire vscode folder
# .vscode/
# Ruff stuff:
.ruff_cache/
# PyPI configuration file
.pypirc
# Marimo
marimo/_static/
marimo/_lsp/
__marimo__/
# Streamlit
.streamlit/secrets.toml
# Temporary files
tmp/
# Docker
docker/
# Hugging Face artifacts
*.safetensors
*.vae
*.pt
# Git (critical for Docker - not in .gitignore)
.git/
.gitignore
.gitattributes
# Documentation
*.md
LICENSE
CHANGELOG.md
docs/
# Tests (commit to git, but exclude from image)
tests/
*_test.py
# IDEs and editors
.vscode/
.idea/
# CI/CD
.github/
# Docker
docker/
Dockerfile
Dockerfile.*
compose.yaml
.dockerignore
# OS
.DS_Store
Thumbs.db
# Development tools config
.pre-commit-config.yaml
.commitizenrc
+5
View File
@@ -0,0 +1,5 @@
# .env.default
# Default cache directory for ComfyUI-Diffusers
# Path is relative to ComfyUI root directory
COMFYUI_DIFFUSERS_CACHE_DIR=./cache
+1
View File
@@ -0,0 +1 @@
COMFYUI_DIFFUSERS_CACHE_DIR=./tmp/cache
-21
View File
@@ -1,21 +0,0 @@
name: Publish to Comfy registry
on:
workflow_dispatch:
push:
branches:
- main
paths:
- "pyproject.toml"
jobs:
publish-node:
name: Publish Custom Node to registry
runs-on: ubuntu-latest
steps:
- name: Check out code
uses: actions/checkout@v4
- name: Publish Custom Node
uses: Comfy-Org/publish-node-action@main
with:
## Add your own personal access token to your Github Repository secrets and reference it here.
personal_access_token: ${{ secrets.REGISTRY_ACCESS_TOKEN }}
+79 -17
View File
@@ -1,6 +1,6 @@
# Byte-compiled / optimized / DLL files
__pycache__/
*.py[cod]
*.py[codz]
*$py.class
# C extensions
@@ -27,8 +27,8 @@ share/python-wheels/
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.
# 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
@@ -46,7 +46,7 @@ htmlcov/
nosetests.xml
coverage.xml
*.cover
*.py,cover
*.py.cover
.hypothesis/
.pytest_cache/
cover/
@@ -92,22 +92,37 @@ ipython_config.py
# 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
# Pipfile.lock
# UV
# Similar to Pipfile.lock, it is generally recommended to include uv.lock in version control.
# This is especially recommended for binary packages to ensure reproducibility, and is more
# commonly ignored for libraries.
# uv.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
# poetry.lock
# poetry.toml
# 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
# pdm recommends including project-wide configuration in pdm.toml, but excluding .pdm-python.
# https://pdm-project.org/en/latest/usage/project/#working-with-version-control
# pdm.lock
# pdm.toml
.pdm-python
.pdm-build/
# pixi
# Similar to Pipfile.lock, it is generally recommended to include pixi.lock in version control.
# pixi.lock
# Pixi creates a virtual environment in the .pixi directory, just like venv module creates one
# in the .venv directory. It is recommended not to include this directory in version control.
.pixi
# PEP 582; used by e.g. github.com/David-OConnor/pyflow and github.com/pdm-project/pdm
__pypackages__/
@@ -116,11 +131,25 @@ __pypackages__/
celerybeat-schedule
celerybeat.pid
# Redis
*.rdb
*.aof
*.pid
# RabbitMQ
mnesia/
rabbitmq/
rabbitmq-data/
# ActiveMQ
activemq-data/
# SageMath parsed files
*.sage.py
# Environments
.env
.envrc
.venv
env/
venv/
@@ -153,10 +182,43 @@ dmypy.json
cython_debug/
# PyCharm
# JetBrains specific template is maintained in a separate JetBrains.gitignore that can
# be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore
# and can be added to the global gitignore or merged into this file. For a more nuclear
# option (not recommended) you can uncomment the following to ignore the entire idea folder.
#.idea/
# JetBrains specific template is maintained in a separate JetBrains.gitignore that can
# be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore
# and can be added to the global gitignore or merged into this file. For a more nuclear
# option (not recommended) you can uncomment the following to ignore the entire idea folder.
# .idea/
StreamDiffusion/*
# Abstra
# Abstra is an AI-powered process automation framework.
# Ignore directories containing user credentials, local state, and settings.
# Learn more at https://abstra.io/docs
.abstra/
# Visual Studio Code
# Visual Studio Code specific template is maintained in a separate VisualStudioCode.gitignore
# that can be found at https://github.com/github/gitignore/blob/main/Global/VisualStudioCode.gitignore
# and can be added to the global gitignore or merged into this file. However, if you prefer,
# you could uncomment the following to ignore the entire vscode folder
# .vscode/
# Ruff stuff:
.ruff_cache/
# PyPI configuration file
.pypirc
# Marimo
marimo/_static/
marimo/_lsp/
__marimo__/
# Streamlit
.streamlit/secrets.toml
# Temporary files
tmp/
# Hugging Face artifacts
*.safetensors
*.vae
*.pt
+235
View File
@@ -0,0 +1,235 @@
{
"cSpell.words": [
"abfs",
"abstra",
"Abstra",
"abstractmethod",
"addopts",
"adlfs",
"aiobotocore",
"aiohttp",
"analyser",
"antlr",
"ANTLR",
"argcomplete",
"argparse",
"armv",
"asyncio",
"authlib",
"autoencoder",
"beautifulsoup",
"bitsandbytes",
"cairocffi",
"celerybeat",
"certifi",
"cfgv",
"Chardet",
"checkdocs",
"cloudpickle",
"codz",
"colorama",
"comfyui",
"commitizen",
"commonmark",
"Connor",
"contextily",
"copybutton",
"cublas",
"CUBLAS",
"cuda",
"CUDA",
"cudnn",
"cufft",
"CUFFT",
"cufile",
"cupti",
"curand",
"CURAND",
"cusolver",
"cusparse",
"CUSPARSE",
"cusparselt",
"cython",
"Cython",
"dask",
"dataframe",
"decli",
"distlib",
"dmypy",
"docstrings",
"dotenv",
"downsample",
"downsamplers",
"dropboxdrivefs",
"dvclive",
"envrc",
"exceptiongroup",
"execnet",
"fastai",
"fastapi",
"fastcore",
"fastparquet",
"filelock",
"filterwarnings",
"flufl",
"fsspec",
"furo",
"fusepy",
"gcsfs",
"gguf",
"gmpy",
"gprof",
"graphviz",
"grimp",
"hdfs",
"htmlcov",
"httpx",
"huggingface",
"idna",
"IDNA",
"igraph",
"importlinter",
"iniconfig",
"intersphinx",
"ipykernel",
"ipynb",
"ipython",
"isort",
"itsdangerous",
"jaraco",
"jaxlib",
"jschema",
"keras",
"kerchunk",
"libcst",
"librosa",
"Limitex",
"linkify",
"loras",
"lxml",
"manylinux",
"Marimo",
"markupsafe",
"matplotlib",
"mdit",
"mdurl",
"mkdocs",
"mnesia",
"momepy",
"moto",
"mpmath",
"musllinux",
"mypy",
"mypyc",
"nccl",
"NCCL",
"networkx",
"nodeenv",
"nosetests",
"numpy",
"numpydoc",
"nvjitlink",
"nvrtc",
"NVRTC",
"nvshmem",
"NVSHMEM",
"nvtx",
"ocifs",
"omegaconf",
"orjson",
"osmnx",
"panflute",
"paramiko",
"pathspec",
"peft",
"pexpect",
"phonemizer",
"pipenv",
"Pipfile",
"psutil",
"pyarrow",
"pybuilder",
"pycache",
"pycodestyle",
"pydantic",
"pydata",
"pydocstyle",
"pydot",
"pyenv",
"pyfakefs",
"pyflakes",
"pyflow",
"pyftpdlib",
"pygit",
"pygraphviz",
"pylint",
"pypa",
"pypackages",
"pypirc",
"pyproject",
"pypy",
"pyrs",
"pytest",
"pytype",
"pyupgrade",
"pyyaml",
"quanto",
"rabbitmq",
"repr",
"rerunfailures",
"resnets",
"riscv",
"ropeproject",
"rtoml",
"safetensors",
"sarif",
"scikit",
"scipy",
"scrapy",
"Scrapy",
"sdist",
"seaborn",
"sentencepiece",
"setuptools",
"simplejson",
"smbprotocol",
"soundfile",
"Spyder",
"spyderproject",
"spyproject",
"stestr",
"streamlit",
"subtests",
"swanlab",
"sympy",
"tensorboard",
"termcolor",
"testpaths",
"testscenarios",
"testtools",
"tiktoken",
"timm",
"tomli",
"tomlkit",
"torchao",
"torchaudio",
"torchdata",
"torchpippy",
"torchvision",
"tqdm",
"trackio",
"typer",
"ukkonen",
"upsamplers",
"urllib",
"venv",
"wandb",
"webassets",
"xarray",
"xdist",
"yarl",
"zarr",
"zipp",
"zstandard"
]
}
-73
View File
@@ -1,74 +1 @@
# ComfyUI-Diffusers
This repository is a custom node in ComfyUI.
## Overview
### Workflow 1
This is a program that allows you to use Hugging Face Diffusers module with ComfyUI. Additionally, Stream Diffusion is also available.
![Workflow](img/workflow.png)
### Workflow 2
In addition, real-time generation is possible by doing the following.
When running, please enable Auto Queue in Extra options.
![Workflow2](img/workflow2.png)
### Workflow 3
In combination with [VideoHelperSuite](https://github.com/Kosinkadink/ComfyUI-VideoHelperSuite.git), you can also run vid2vid.
![Workflow3](img/workflow3.png)
The execution looks like this:
![Workflow3demo](img/workflow3demo.gif)
## Usage
Run the following command inside ComfyUI/custom_nodes.
```cmd
git clone https://github.com/Limitex/ComfyUI-Diffusers.git
cd ComfyUI-Diffusers
pip install -r requirements.txt
git clone https://github.com/cumulo-autumn/StreamDiffusion.git
python -m streamdiffusion.tools.install-tensorrt
```
## Recommended Custom Nodes
https://github.com/Kosinkadink/ComfyUI-VideoHelperSuite.git
## Node
### Diffusers Pipeline Loader (DiffusersPipelineLoader)
### Diffusers Vae Loader (DiffusersVaeLoader)
### Diffusers Scheduler Loader (DiffusersSchedulerLoader)
### Diffusers Model Makeup (DiffusersModelMakeup)
### Diffusers Clip Text Encode (DiffusersClipTextEncode)
### Diffusers Sampler (DiffusersSampler)
### Create Int List (CreateIntListNode)
### LcmLoraLoader (LcmLoraLoader)
### StreamDiffusion Create Stream (StreamDiffusionCreateStream)
### StreamDiffusion Sampler (StreamDiffusionSampler)
### StreamDiffusion Warmup (StreamDiffusionWarmup)
### StreamDiffusion Fast Sampler (StreamDiffusionFastSampler)
## Reference
https://github.com/cumulo-autumn/StreamDiffusion
+175
View File
@@ -0,0 +1,175 @@
version: '3'
vars:
PYTHON_SRC: "src"
PYTHON_TESTS: "tests"
LINT_TARGETS: "src tests"
tasks:
default:
desc: "List all available tasks"
cmds:
- task --list
install:
desc: "Install all dependencies (including dev) based on poetry.lock"
cmds:
- poetry install
- task: export
install:prod:
desc: "Install only production dependencies"
cmds:
- poetry install --without dev
- task: export
update:
desc: "Update dependencies based on pyproject.toml and regenerate poetry.lock"
cmds:
- poetry update
- task: export
add:
desc: "Add a new production dependency (usage: task add -- requests)"
cmds:
- poetry add {{.CLI_ARGS}}
- task: export
preconditions:
- sh: "[ -n '{{.CLI_ARGS}}' ]"
msg: "CLI_ARGS must be set (e.g., task add -- requests)"
add:dev:
desc: "Add a new development dependency (usage: task add:dev -- pytest)"
cmds:
- poetry add {{.CLI_ARGS}} --group dev
- task: export
preconditions:
- sh: "[ -n '{{.CLI_ARGS}}' ]"
msg: "CLI_ARGS must be set (e.g., task add:dev -- pytest)"
remove:
desc: "Remove a dependency (usage: task remove -- requests)"
cmds:
- poetry remove {{.CLI_ARGS}}
- task: export
preconditions:
- sh: "[ -n '{{.CLI_ARGS}}' ]"
msg: "CLI_ARGS must be set (e.g., task remove -- requests)"
shell:
desc: "Activate the virtual environment shell managed by Poetry"
cmds:
- poetry shell
run:
desc: "Run a command within the virtual environment (usage: task run -- python main.py)"
cmds:
- poetry run {{.CLI_ARGS}}
preconditions:
- sh: "[ -n '{{.CLI_ARGS}}' ]"
msg: "CLI_ARGS must be set (e.g., task run -- python main.py)"
format:
desc: "Automatically format code using Ruff"
cmds:
- poetry run ruff format {{.LINT_TARGETS}}
format:check:
desc: "Check code formatting without making changes"
cmds:
- poetry run ruff format {{.LINT_TARGETS}} --check
lint:
desc: "Run linting with ruff"
cmds:
- poetry run ruff check {{.LINT_TARGETS}}
lint:fix:
desc: "Run linting with auto-fix"
cmds:
- poetry run ruff check {{.LINT_TARGETS}} --fix
type:
desc: "Run static type checking using mypy"
cmds:
- poetry run mypy {{.PYTHON_SRC}}
fix:
desc: "Auto-fix all formatting and linting issues"
cmds:
- task: format
- task: lint:fix
check:
desc: "Run all checks (format check, lint, type check, and tests)"
cmds:
- task: format:check
- task: lint
- task: type
- task: test
test:
desc: "Run tests using pytest"
cmds:
- poetry run pytest {{.PYTHON_TESTS}}
test:verbose:
desc: "Run tests with verbose output"
cmds:
- poetry run pytest -v {{.PYTHON_TESTS}}
test:cov:
desc: "Run tests with coverage report"
cmds:
- poetry run pytest --cov={{.PYTHON_SRC}} {{.PYTHON_TESTS}} --cov-report=html --cov-report=term
build:
desc: "Build sdist and wheel"
cmds:
- poetry build
export:
desc: "Export production dependencies to requirements.txt"
cmds:
- poetry export -f requirements.txt --output requirements.txt --without-hashes
export:dev:
desc: "Export all dependencies (including dev) to requirements-dev.txt"
cmds:
- poetry export -f requirements.txt --output requirements-dev.txt --without-hashes --with dev
export:all:
desc: "Export both prod and dev requirements"
cmds:
- task: export
- task: export:dev
lookup:
desc: "Check available versions of a package on PyPI (usage: task lookup -- requests)"
cmds:
- poetry run pip index versions {{.CLI_ARGS}}
preconditions:
- sh: "[ -n '{{.CLI_ARGS}}' ]"
msg: "CLI_ARGS must be set (e.g., task lookup -- requests)"
tree:
desc: "Show dependency tree of a package (usage: task tree -- requests)"
cmds:
- poetry show {{.CLI_ARGS}} --tree
preconditions:
- sh: "[ -n '{{.CLI_ARGS}}' ]"
msg: "CLI_ARGS must be set (e.g., task tree -- requests)"
clean:
desc: "Clean up cache and temporary files"
cmds:
- rm -rf .pytest_cache
- rm -rf .mypy_cache
- rm -rf .ruff_cache
- rm -rf htmlcov
- rm -rf .coverage
- rm -rf dist
- rm -rf build
- find . -type d -name __pycache__ -exec rm -rf {} +
- find . -type d -name "*.egg-info" -exec rm -rf {} +
- find . -type f -name "*.pyc" -delete
+21 -2
View File
@@ -1,3 +1,22 @@
from .nodes import *
from .src.config import load_envs
from .src.nodes import (
NODE_CLASS_MAPPINGS,
NODE_DISPLAY_NAME_MAPPINGS,
ComfyUIAutoencoderDTO,
ComfyUIClipDTO,
ComfyUIConditioningDTO,
ComfyUIPipelineDTO,
ComfyUISchedulerDTO,
)
__all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS"]
load_envs()
__all__ = [
"NODE_CLASS_MAPPINGS",
"NODE_DISPLAY_NAME_MAPPINGS",
"ComfyUIPipelineDTO",
"ComfyUIAutoencoderDTO",
"ComfyUIClipDTO",
"ComfyUIConditioningDTO",
"ComfyUISchedulerDTO",
]
+38
View File
@@ -0,0 +1,38 @@
services:
comfyui:
build:
context: .
dockerfile: docker/comfyui/Dockerfile
container_name: comfyui
runtime: nvidia
environment:
- CUDA_VISIBLE_DEVICES=0
- NVIDIA_VISIBLE_DEVICES=all
- NVIDIA_DRIVER_CAPABILITIES=compute,utility
- POETRY_VENV_IN_PROJECT=false
- POETRY_NO_INTERACTION=1
- POETRY_CACHE_DIR=/tmp/poetry_cache
volumes:
- ./tmp/models/checkpoints:/app/ComfyUI/models/checkpoints
- ./tmp/models/loras:/app/ComfyUI/models/loras
- ./tmp/models/vae:/app/ComfyUI/models/vae
- ./tmp/output:/app/ComfyUI/output
- ./tmp/workflows:/app/ComfyUI/user/default/workflows
- ./tmp/tmp:/app/ComfyUI/temp
- .:/app/ComfyUI/custom_nodes/ComfyUI-Diffusers
ports:
- "8188:8188"
restart: unless-stopped
stdin_open: true
tty: true
dev:
build:
context: .
dockerfile: docker/dev/Dockerfile
volumes:
- .:/app
- venv_data:/cache/.venv
volumes:
venv_data:
+22
View File
@@ -0,0 +1,22 @@
FROM alpine/git:latest AS builder
WORKDIR /app
RUN git clone --depth 1 --branch v0.3.67 https://github.com/comfyanonymous/ComfyUI.git ComfyUI && \
rm -rf ComfyUI/.git
FROM pytorch/pytorch:2.6.0-cuda12.4-cudnn9-runtime
WORKDIR /app
COPY --from=builder /app/ComfyUI /app/ComfyUI
WORKDIR /app/ComfyUI
RUN apt-get update && \
apt-get install -y --no-install-recommends \
git \
curl \
build-essential \
&& pip install --no-cache-dir -r requirements.txt \
&& apt-get remove -y build-essential curl \
&& apt-get autoremove -y && \
rm -rf /var/lib/apt/lists/*
COPY ./requirements.txt ./tmp.txt
RUN pip install --no-cache-dir -r tmp.txt && rm tmp.txt
EXPOSE 8188
CMD ["python", "main.py", "--listen", "0.0.0.0", "--port", "8188"]
+24
View File
@@ -0,0 +1,24 @@
FROM python:3.11-slim-bookworm
ENV POETRY_HOME="/opt/poetry"
ENV POETRY_VERSION="2.2.1"
ENV POETRY_VIRTUALENVS_PATH="/cache/.venv"
ENV POETRY_VIRTUALENVS_IN_PROJECT="false"
ENV PATH="$POETRY_HOME/bin:$PATH"
WORKDIR /app
RUN mkdir -p /cache/.venv && chmod 777 /cache/.venv
RUN apt-get update && \
apt-get install -y --no-install-recommends \
git \
curl \
build-essential \
ca-certificates \
&& curl -sSL https://install.python-poetry.org | python - \
&& curl -1sLf 'https://dl.cloudsmith.io/public/task/task/setup.deb.sh' | bash \
&& apt-get update \
&& apt-get install -y task \
&& apt-get clean \
&& rm -rf /var/lib/apt/lists/*
RUN poetry self add \
poetry-plugin-export \
poetry-plugin-shell
RUN task --version && poetry --version
BIN
View File
Binary file not shown.

Before

Width:  |  Height:  |  Size: 1.1 MiB

BIN
View File
Binary file not shown.

Before

Width:  |  Height:  |  Size: 476 KiB

BIN
View File
Binary file not shown.

Before

Width:  |  Height:  |  Size: 339 KiB

Binary file not shown.

Before

Width:  |  Height:  |  Size: 9.0 MiB

+86
View File
@@ -0,0 +1,86 @@
[mypy]
python_version = 3.10
namespace_packages = True
explicit_package_bases = True
strict = True
warn_return_any = True
warn_unused_configs = True
disallow_untyped_defs = True
disallow_incomplete_defs = True
check_untyped_defs = True
no_implicit_optional = True
warn_redundant_casts = True
warn_unused_ignores = True
warn_no_return = True
warn_unreachable = True
strict_equality = True
# Main module
[mypy-__main__.*]
ignore_missing_imports = True
[mypy-comfy.*]
ignore_missing_imports = True
# ComfyUI folder_paths module
[mypy-folder_paths]
ignore_missing_imports = True
[mypy-folder_paths.*]
ignore_missing_imports = True
# Diffusers library
[mypy-diffusers]
ignore_missing_imports = True
[mypy-diffusers.*]
ignore_missing_imports = True
# Dependency Injector
[mypy-dependency_injector]
ignore_missing_imports = True
[mypy-dependency_injector.*]
ignore_missing_imports = True
# Transformers library
[mypy-transformers]
ignore_missing_imports = True
[mypy-transformers.*]
ignore_missing_imports = True
# PyTorch
[mypy-torch]
ignore_missing_imports = True
[mypy-torch.*]
ignore_missing_imports = True
# TorchVision
[mypy-torchvision]
ignore_missing_imports = True
[mypy-torchvision.*]
ignore_missing_imports = True
# TorchAudio
[mypy-torchaudio]
ignore_missing_imports = True
[mypy-torchaudio.*]
ignore_missing_imports = True
# OmegaConf
[mypy-omegaconf]
ignore_missing_imports = True
[mypy-omegaconf.*]
ignore_missing_imports = True
# SafeTensors
[mypy-safetensors]
ignore_missing_imports = True
[mypy-safetensors.*]
ignore_missing_imports = True
-435
View File
@@ -1,435 +0,0 @@
import copy
import os
import torch
from safetensors.torch import load_file
from .utils import SCHEDULERS, token_auto_concat_embeds, vae_pt_to_vae_diffuser, convert_images_to_tensors, convert_tensors_to_images, resize_images
from comfy.model_management import get_torch_device
import folder_paths
from streamdiffusion import StreamDiffusion
from streamdiffusion.image_utils import postprocess_image
from diffusers import StableDiffusionPipeline, AutoencoderKL, AutoencoderTiny
class DiffusersPipelineLoader:
def __init__(self):
self.tmp_dir = folder_paths.get_temp_directory()
self.dtype = torch.float32
@classmethod
def INPUT_TYPES(s):
return {"required": { "ckpt_name": (folder_paths.get_filename_list("checkpoints"), ), }}
RETURN_TYPES = ("PIPELINE", "AUTOENCODER", "SCHEDULER",)
FUNCTION = "create_pipeline"
CATEGORY = "Diffusers"
def create_pipeline(self, ckpt_name):
ckpt_cache_path = os.path.join(self.tmp_dir, ckpt_name)
StableDiffusionPipeline.from_single_file(
pretrained_model_link_or_path=folder_paths.get_full_path("checkpoints", ckpt_name),
torch_dtype=self.dtype,
cache_dir=self.tmp_dir,
).save_pretrained(ckpt_cache_path, safe_serialization=True)
pipe = StableDiffusionPipeline.from_pretrained(
pretrained_model_name_or_path=ckpt_cache_path,
torch_dtype=self.dtype,
cache_dir=self.tmp_dir,
)
return ((pipe, ckpt_cache_path), pipe.vae, pipe.scheduler)
class DiffusersVaeLoader:
def __init__(self):
self.tmp_dir = folder_paths.get_temp_directory()
self.dtype = torch.float32
@classmethod
def INPUT_TYPES(s):
return {"required": { "vae_name": (folder_paths.get_filename_list("vae"), ), }}
RETURN_TYPES = ("AUTOENCODER",)
FUNCTION = "create_pipeline"
CATEGORY = "Diffusers"
def create_pipeline(self, vae_name):
ckpt_cache_path = os.path.join(self.tmp_dir, vae_name)
vae_pt_to_vae_diffuser(folder_paths.get_full_path("vae", vae_name), ckpt_cache_path)
vae = AutoencoderKL.from_pretrained(
pretrained_model_name_or_path=ckpt_cache_path,
torch_dtype=self.dtype,
cache_dir=self.tmp_dir,
)
return (vae,)
class DiffusersSchedulerLoader:
def __init__(self):
self.tmp_dir = folder_paths.get_temp_directory()
self.dtype = torch.float32
@classmethod
def INPUT_TYPES(s):
return {
"required": {
"pipeline": ("PIPELINE", ),
"scheduler_name": (list(SCHEDULERS.keys()), ),
}
}
RETURN_TYPES = ("SCHEDULER",)
FUNCTION = "load_scheduler"
CATEGORY = "Diffusers"
def load_scheduler(self, pipeline, scheduler_name):
scheduler = SCHEDULERS[scheduler_name].from_pretrained(
pretrained_model_name_or_path=pipeline[1],
torch_dtype=self.dtype,
cache_dir=self.tmp_dir,
subfolder='scheduler'
)
return (scheduler,)
class DiffusersModelMakeup:
def __init__(self):
self.torch_device = get_torch_device()
@classmethod
def INPUT_TYPES(s):
return {
"required": {
"pipeline": ("PIPELINE", ),
"scheduler": ("SCHEDULER", ),
"autoencoder": ("AUTOENCODER", ),
},
}
RETURN_TYPES = ("MAKED_PIPELINE",)
FUNCTION = "makeup_pipeline"
CATEGORY = "Diffusers"
def makeup_pipeline(self, pipeline, scheduler, autoencoder):
pipeline = pipeline[0]
pipeline.vae = autoencoder
pipeline.scheduler = scheduler
pipeline.safety_checker = None if pipeline.safety_checker is None else lambda images, **kwargs: (images, [False])
pipeline.enable_attention_slicing()
pipeline = pipeline.to(self.torch_device)
return (pipeline,)
class DiffusersClipTextEncode:
@classmethod
def INPUT_TYPES(s):
return {"required": {
"maked_pipeline": ("MAKED_PIPELINE", ),
"positive": ("STRING", {"multiline": True}),
"negative": ("STRING", {"multiline": True}),
}}
RETURN_TYPES = ("EMBEDS", "EMBEDS", "STRING", "STRING", )
RETURN_NAMES = ("positive_embeds", "negative_embeds", "positive", "negative", )
FUNCTION = "concat_embeds"
CATEGORY = "Diffusers"
def concat_embeds(self, maked_pipeline, positive, negative):
positive_embeds, negative_embeds = token_auto_concat_embeds(maked_pipeline, positive,negative)
return (positive_embeds, negative_embeds, positive, negative, )
class DiffusersSampler:
def __init__(self):
self.torch_device = get_torch_device()
@classmethod
def INPUT_TYPES(s):
return {"required": {
"maked_pipeline": ("MAKED_PIPELINE", ),
"positive_embeds": ("EMBEDS", ),
"negative_embeds": ("EMBEDS", ),
"width": ("INT", {"default": 512, "min": 1, "max": 8192, "step": 1}),
"height": ("INT", {"default": 512, "min": 1, "max": 8192, "step": 1}),
"steps": ("INT", {"default": 20, "min": 1, "max": 10000}),
"cfg": ("FLOAT", {"default": 8.0, "min": 0.0, "max": 100.0, "step":0.1, "round": 0.01}),
"seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff}),
}}
RETURN_TYPES = ("IMAGE",)
FUNCTION = "sample"
CATEGORY = "Diffusers"
def sample(self, maked_pipeline, positive_embeds, negative_embeds, height, width, steps, cfg, seed):
images = maked_pipeline(
prompt_embeds=positive_embeds,
height=height,
width=width,
num_inference_steps=steps,
guidance_scale=cfg,
negative_prompt_embeds=negative_embeds,
generator=torch.Generator(self.torch_device).manual_seed(seed)
).images
return (convert_images_to_tensors(images),)
# - Stream Diffusion -
class CreateIntListNode:
@classmethod
def INPUT_TYPES(s):
max_element = 10
return {
"required": {
"elements_count" : ("INT", {"default": 2, "min": 1, "max": max_element, "step": 1}),
},
"optional": {
f"element_{i}": ("INT", {"default": 0}) for i in range(1, max_element)
}
}
RETURN_TYPES = ("LIST",)
FUNCTION = "create_list"
CATEGORY = "Diffusers/StreamDiffusion"
def create_list(self, elements_count, **kwargs):
return ([value for key, value in kwargs.items()][:elements_count], )
class LcmLoraLoader:
@classmethod
def INPUT_TYPES(s):
return {"required": { "lora_name": (folder_paths.get_filename_list("loras"), ), }}
RETURN_TYPES = ("LCM_LORA",)
FUNCTION = "load_lora"
CATEGORY = "Diffusers/StreamDiffusion"
def load_lora(self, lora_name):
return (load_file(folder_paths.get_full_path("loras", lora_name)), )
class StreamDiffusionCreateStream:
def __init__(self):
self.dtype = torch.float32
self.torch_device = get_torch_device()
self.tmp_dir = folder_paths.get_temp_directory()
@classmethod
def INPUT_TYPES(s):
return {
"required": {
"maked_pipeline": ("MAKED_PIPELINE", ),
"t_index_list": ("LIST", ),
"width": ("INT", {"default": 512, "min": 1, "max": 8192, "step": 1}),
"height": ("INT", {"default": 512, "min": 1, "max": 8192, "step": 1}),
"do_add_noise": ("BOOLEAN", {"default": True}),
"use_denoising_batch": ("BOOLEAN", {"default": True}),
"frame_buffer_size": ("INT", {"default": 1, "min": 1, "max": 10000}),
"cfg_type": (["none", "full", "self", "initialize"], {"default": "none"}),
"xformers_memory_efficient_attention": ("BOOLEAN", {"default": False}),
"lcm_lora" : ("LCM_LORA", ),
"tiny_vae" : ("STRING", {"default": "madebyollin/taesd"})
},
}
RETURN_TYPES = ("STREAM",)
FUNCTION = "load_stream"
CATEGORY = "Diffusers/StreamDiffusion"
def load_stream(self, maked_pipeline, t_index_list, width, height, do_add_noise, use_denoising_batch, frame_buffer_size, cfg_type, xformers_memory_efficient_attention, lcm_lora, tiny_vae):
maked_pipeline = copy.deepcopy(maked_pipeline)
lcm_lora = copy.deepcopy(lcm_lora)
stream = StreamDiffusion(
pipe = maked_pipeline,
t_index_list = t_index_list,
torch_dtype = self.dtype,
width = width,
height = height,
do_add_noise = do_add_noise,
use_denoising_batch = use_denoising_batch,
frame_buffer_size = frame_buffer_size,
cfg_type = cfg_type,
)
stream.load_lcm_lora(lcm_lora)
stream.fuse_lora()
stream.vae = AutoencoderTiny.from_pretrained(
pretrained_model_name_or_path=tiny_vae,
torch_dtype=self.dtype,
cache_dir=self.tmp_dir,
).to(
device=maked_pipeline.device,
dtype=maked_pipeline.dtype
)
if xformers_memory_efficient_attention:
maked_pipeline.enable_xformers_memory_efficient_attention()
return (stream, )
class StreamDiffusionSampler:
def __init__(self):
self.torch_device = get_torch_device()
@classmethod
def INPUT_TYPES(s):
return {
"required": {
"stream": ("STREAM", ),
"positive": ("STRING", {"multiline": True}),
"negative": ("STRING", {"multiline": True}),
"steps": ("INT", {"default": 50, "min": 1, "max": 10000}),
"cfg": ("FLOAT", {"default": 1.2, "min": 0.0, "max": 100.0}),
"delta": ("FLOAT", {"default": 1, "min": 0.0, "max": 1.0}),
"seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff}),
"num": ("INT", {"default": 1, "min": 1, "max": 10000}),
"warmup": ("INT", {"default": 1, "min": 0, "max": 10000}),
},
"optional" : {
"image" : ("IMAGE", )
}
}
RETURN_TYPES = ("IMAGE",)
FUNCTION = "sample"
CATEGORY = "Diffusers/StreamDiffusion"
def sample(self, stream: StreamDiffusion, positive, negative, steps, cfg, delta, seed, num, warmup, image = None):
stream.prepare(
prompt = positive,
negative_prompt = negative,
num_inference_steps = steps,
guidance_scale = cfg,
delta = delta,
seed = seed
)
if image != None:
image = convert_tensors_to_images(image)
image = resize_images(image, (stream.width, stream.height))
for _ in range(warmup):
stream()
result = []
for _ in range(num):
x_outputs = []
if image is None:
x_outputs.append(stream.txt2img())
else:
stream(image[0])
for i in image[1:] + image[-1:]:
x_outputs.append(stream(i))
for x_output in x_outputs:
result.append(postprocess_image(x_output, output_type="pil")[0])
return (convert_images_to_tensors(result),)
class StreamDiffusionWarmup:
@classmethod
def INPUT_TYPES(s):
return {
"required": {
"stream": ("STREAM", ),
"negative": ("STRING", {"multiline": True}),
"steps": ("INT", {"default": 50, "min": 1, "max": 10000}),
"cfg": ("FLOAT", {"default": 1.2, "min": 0.0, "max": 100.0}),
"delta": ("FLOAT", {"default": 1, "min": 0.0, "max": 1.0}),
"seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff}),
"warmup": ("INT", {"default": 1, "min": 0, "max": 10000}),
},
}
RETURN_TYPES = ("WARMUP_STREAM",)
FUNCTION = "stream_warmup"
CATEGORY = "Diffusers/StreamDiffusion"
def stream_warmup(self, stream: StreamDiffusion, negative, steps, cfg, delta, seed, warmup):
stream.prepare(
prompt="",
negative_prompt=negative,
num_inference_steps = steps,
guidance_scale = cfg,
delta = delta,
seed = seed
)
for _ in range(warmup):
stream()
return (stream, )
class StreamDiffusionFastSampler:
@classmethod
def INPUT_TYPES(s):
return {
"required": {
"warmup_stream": ("WARMUP_STREAM", ),
"positive": ("STRING", {"multiline": True}),
"num": ("INT", {"default": 1, "min": 1, "max": 10000}),
},
}
RETURN_TYPES = ("IMAGE",)
FUNCTION = "sample"
CATEGORY = "Diffusers/StreamDiffusion"
def sample(self, warmup_stream, positive, num):
stream: StreamDiffusion = warmup_stream
stream.update_prompt(positive)
result = []
for _ in range(num):
x_output = stream.txt2img()
result.append(postprocess_image(x_output, output_type="pil")[0])
return (convert_images_to_tensors(result),)
# - - - - - - - - - - - - - - - - - -
NODE_CLASS_MAPPINGS = {
"DiffusersPipelineLoader": DiffusersPipelineLoader,
"DiffusersVaeLoader": DiffusersVaeLoader,
"DiffusersSchedulerLoader": DiffusersSchedulerLoader,
"DiffusersModelMakeup": DiffusersModelMakeup,
"DiffusersClipTextEncode": DiffusersClipTextEncode,
"DiffusersSampler": DiffusersSampler,
"CreateIntListNode": CreateIntListNode,
"LcmLoraLoader": LcmLoraLoader,
"StreamDiffusionCreateStream": StreamDiffusionCreateStream,
"StreamDiffusionSampler": StreamDiffusionSampler,
"StreamDiffusionWarmup": StreamDiffusionWarmup,
"StreamDiffusionFastSampler": StreamDiffusionFastSampler,
}
NODE_DISPLAY_NAME_MAPPINGS = {
"DiffusersPipelineLoader": "Diffusers Pipeline Loader",
"DiffusersVaeLoader": "Diffusers Vae Loader",
"DiffusersSchedulerLoader": "Diffusers Scheduler Loader",
"DiffusersModelMakeup": "Diffusers Model Makeup",
"DiffusersClipTextEncode": "Diffusers Clip Text Encode",
"DiffusersSampler": "Diffusers Sampler",
"CreateIntListNode": "Create Int List",
"LcmLoraLoader": "LCM Lora Loader",
"StreamDiffusionCreateStream": "StreamDiffusion Create Stream",
"StreamDiffusionSampler": "StreamDiffusion Sampler",
"StreamDiffusionWarmup": "StreamDiffusion Warmup",
"StreamDiffusionFastSampler": "StreamDiffusion Fast Sampler",
}
Generated
+2950
View File
File diff suppressed because it is too large Load Diff
+186 -12
View File
@@ -1,15 +1,189 @@
[project]
name = "comfyui-diffusers"
description = "This extension enables the use of the diffuser pipeline in ComfyUI. It also includes nodes related to Stream Diffusion."
version = "1.0.0"
license = "LICENSE"
dependencies = ["diffusers[torch]", "accelerate", "transformers", "safetensors", "omegaconf", "pytorch_lightning", "xformers", "git+https://github.com/cumulo-autumn/StreamDiffusion.git@main#egg=streamdiffusion[tensorrt]"]
name = "ComfyUI-Diffusers"
version = "0.1.0"
description = "This repository is a custom node in ComfyUI. This is a program that allows you to use Huggingface Diffusers module with ComfyUI. Additionally, Stream Diffusion is also available."
authors = [
{name = "Limitex", email = "76650151+Limitex@users.noreply.github.com"}
]
readme = "README.md"
requires-python = ">=3.10,<3.15"
dependencies = [
"dependency-injector (>=4.48.2,<5.0.0)",
"diffusers[torch] (>=0.35.2,<0.36.0)",
"torch (>=2.9.0,<3.0.0)",
"torchaudio (>=2.9.0,<3.0.0)",
"torchvision (>=0.24.0,<0.25.0)",
"transformers (>=4.57.1,<5.0.0)",
"omegaconf (>=2.3.0,<3.0.0)",
"safetensors (>=0.4.0,<0.5.0)",
"requests (>=2.32.0,<3.0.0)",
"python-dotenv (>=1.2.1,<2.0.0)",
]
[project.urls]
Repository = "https://github.com/Limitex/ComfyUI-Diffusers"
# Used by Comfy Registry https://comfyregistry.org
[tool.poetry]
packages = [
{include = "src"}
]
[tool.comfy]
PublisherId = "limitex"
DisplayName = "ComfyUI-Diffusers"
Icon = ""
[tool.pytest.ini_options]
testpaths = ["tests"]
python_files = ["test_*.py", "*_test.py"]
python_classes = ["Test*"]
python_functions = ["test_*"]
addopts = [
"--strict-markers",
"--strict-config",
"--cov=src",
"--cov-branch",
"--cov-report=term-missing:skip-covered",
"--cov-report=html",
"--cov-report=xml",
"--cov-fail-under=80",
"-v",
"-ra",
]
markers = [
"unit: Unit tests",
"integration: Integration tests",
"slow: Slow tests",
]
filterwarnings = [
"error",
"ignore::UserWarning",
"ignore::DeprecationWarning",
]
[tool.coverage.run]
source = ["src"]
omit = [
"*/tests/*",
"*/test_*.py",
"*/__init__.py",
"src/di/*", # DI container doesn't need test coverage
]
branch = true
parallel = true
[tool.coverage.report]
precision = 2
show_missing = true
skip_covered = false
exclude_lines = [
"pragma: no cover",
"def __repr__",
"def __str__",
"raise AssertionError",
"raise NotImplementedError",
"if __name__ == .__main__.:",
"if TYPE_CHECKING:",
"if typing.TYPE_CHECKING:",
"@abstractmethod",
"@abc.abstractmethod",
"@overload",
"class .*\\bProtocol\\):",
"pass",
]
[tool.coverage.html]
directory = "htmlcov"
[tool.coverage.xml]
output = "coverage.xml"
[tool.bandit]
exclude_dirs = ["tests", "venv", ".venv", ".tox", "build", "dist"]
tests = [
"B201", "B301", "B302", "B303", "B304", "B305", "B306", "B307",
"B308", "B309", "B310", "B311", "B312", "B313", "B314", "B315",
"B316", "B317", "B318", "B319", "B320", "B321", "B323", "B324", "B325"
]
skips = ["B101", "B601"] # Allow assert in tests & shell=True if needed
[tool.bandit.assert_used]
skips = ["*/test_*.py", "*/*_test.py"]
[tool.importlinter]
root_package = "src"
[[tool.importlinter.contracts]]
name = "Clean Architecture Layers"
type = "layers"
layers = [
"src.nodes", # Presentation Layer (ComfyUI Nodes)
"src.ui", # Presentation Layer (Handlers)
"src.service", # Application Layer (Use Cases)
"src.domain", # Domain Layer (Entities & Interfaces)
]
containers = ["src"]
[[tool.importlinter.contracts]]
name = "Domain Layer Independence"
type = "forbidden"
source_modules = ["src.domain"]
forbidden_modules = [
"src.service",
"src.ui",
"src.nodes",
"src.infra",
"src.di",
]
[[tool.importlinter.contracts]]
name = "Application Layer Independence"
type = "forbidden"
source_modules = ["src.service"]
forbidden_modules = [
"src.ui",
"src.nodes",
"src.infra",
"src.di",
]
[[tool.importlinter.contracts]]
name = "Infrastructure depends only on Domain"
type = "forbidden"
source_modules = ["src.infra"]
forbidden_modules = [
"src.service",
"src.ui",
"src.nodes",
"src.di",
]
[[tool.importlinter.contracts]]
name = "Presentation Layer can use Application & Domain"
type = "forbidden"
source_modules = ["src.ui", "src.nodes"]
forbidden_modules = [
"src.infra",
]
[tool.commitizen]
name = "cz_conventional_commits"
version = "0.1.0"
tag_format = "v$version"
version_files = [
"pyproject.toml:version"
]
update_changelog_on_bump = true
changelog_file = "CHANGELOG.md"
[build-system]
requires = ["poetry-core>=2.0.0,<3.0.0"]
build-backend = "poetry.core.masonry.api"
[dependency-groups]
dev = [
"ruff (>=0.14.2,<0.15.0)",
"mypy (>=1.18.2,<2.0.0)",
"types-requests (>=2.32.4.20250913,<3.0.0.0)",
"pytest (>=8.4.2,<9.0.0)",
"pytest-cov (>=7.0.0,<8.0.0)",
"pytest-asyncio (>=1.2.0,<2.0.0)",
"pytest-mock (>=3.15.1,<4.0.0)",
"pytest-xdist (>=3.8.0,<4.0.0)",
"bandit[toml] (>=1.8.6,<2.0.0)",
"import-linter (>=2.5.2,<3.0.0)",
"pre-commit (>=4.3.0,<5.0.0)",
"commitizen (>=4.9.1,<5.0.0)",
]
+54 -8
View File
@@ -1,8 +1,54 @@
diffusers[torch]
accelerate
transformers
safetensors
omegaconf
pytorch_lightning
xformers
git+https://github.com/cumulo-autumn/StreamDiffusion.git@main#egg=streamdiffusion[tensorrt]
accelerate==1.11.0 ; python_version >= "3.10" and python_version < "3.15"
antlr4-python3-runtime==4.9.3 ; python_version >= "3.10" and python_version < "3.15"
certifi==2025.10.5 ; python_version >= "3.10" and python_version < "3.15"
charset-normalizer==3.4.4 ; python_version >= "3.10" and python_version < "3.15"
colorama==0.4.6 ; python_version >= "3.10" and python_version < "3.15" and platform_system == "Windows"
dependency-injector==4.48.2 ; python_version >= "3.10" and python_version < "3.15"
diffusers==0.35.2 ; python_version >= "3.10" and python_version < "3.15"
filelock==3.20.0 ; python_version >= "3.10" and python_version < "3.15"
fsspec==2025.10.0 ; python_version >= "3.10" and python_version < "3.15"
hf-xet==1.2.0 ; python_version >= "3.10" and python_version < "3.15" and (platform_machine == "x86_64" or platform_machine == "amd64" or platform_machine == "arm64" or platform_machine == "aarch64")
huggingface-hub==0.36.0 ; python_version >= "3.10" and python_version < "3.15"
idna==3.11 ; python_version >= "3.10" and python_version < "3.15"
importlib-metadata==8.7.0 ; python_version >= "3.10" and python_version < "3.15"
jinja2==3.1.6 ; python_version >= "3.10" and python_version < "3.15"
markupsafe==3.0.3 ; python_version >= "3.10" and python_version < "3.15"
mpmath==1.3.0 ; python_version >= "3.10" and python_version < "3.15"
networkx==3.4.2 ; python_version >= "3.10" and python_version < "3.15"
numpy==2.2.6 ; python_version >= "3.10" and python_version < "3.15"
nvidia-cublas-cu12==12.8.4.1 ; python_version >= "3.10" and python_version < "3.15" and platform_system == "Linux" and platform_machine == "x86_64"
nvidia-cuda-cupti-cu12==12.8.90 ; python_version >= "3.10" and python_version < "3.15" and platform_system == "Linux" and platform_machine == "x86_64"
nvidia-cuda-nvrtc-cu12==12.8.93 ; python_version >= "3.10" and python_version < "3.15" and platform_system == "Linux" and platform_machine == "x86_64"
nvidia-cuda-runtime-cu12==12.8.90 ; python_version >= "3.10" and python_version < "3.15" and platform_system == "Linux" and platform_machine == "x86_64"
nvidia-cudnn-cu12==9.10.2.21 ; python_version >= "3.10" and python_version < "3.15" and platform_system == "Linux" and platform_machine == "x86_64"
nvidia-cufft-cu12==11.3.3.83 ; python_version >= "3.10" and python_version < "3.15" and platform_system == "Linux" and platform_machine == "x86_64"
nvidia-cufile-cu12==1.13.1.3 ; python_version >= "3.10" and python_version < "3.15" and platform_system == "Linux" and platform_machine == "x86_64"
nvidia-curand-cu12==10.3.9.90 ; python_version >= "3.10" and python_version < "3.15" and platform_system == "Linux" and platform_machine == "x86_64"
nvidia-cusolver-cu12==11.7.3.90 ; python_version >= "3.10" and python_version < "3.15" and platform_system == "Linux" and platform_machine == "x86_64"
nvidia-cusparse-cu12==12.5.8.93 ; python_version >= "3.10" and python_version < "3.15" and platform_system == "Linux" and platform_machine == "x86_64"
nvidia-cusparselt-cu12==0.7.1 ; python_version >= "3.10" and python_version < "3.15" and platform_system == "Linux" and platform_machine == "x86_64"
nvidia-nccl-cu12==2.27.5 ; python_version >= "3.10" and python_version < "3.15" and platform_system == "Linux" and platform_machine == "x86_64"
nvidia-nvjitlink-cu12==12.8.93 ; python_version >= "3.10" and python_version < "3.15" and platform_system == "Linux" and platform_machine == "x86_64"
nvidia-nvshmem-cu12==3.3.20 ; python_version >= "3.10" and python_version < "3.15" and platform_system == "Linux" and platform_machine == "x86_64"
nvidia-nvtx-cu12==12.8.90 ; python_version >= "3.10" and python_version < "3.15" and platform_system == "Linux" and platform_machine == "x86_64"
omegaconf==2.3.0 ; python_version >= "3.10" and python_version < "3.15"
packaging==25.0 ; python_version >= "3.10" and python_version < "3.15"
pillow==12.0.0 ; python_version >= "3.10" and python_version < "3.15"
psutil==7.1.2 ; python_version >= "3.10" and python_version < "3.15"
python-dotenv==1.2.1 ; python_version >= "3.10" and python_version < "3.15"
pyyaml==6.0.3 ; python_version >= "3.10" and python_version < "3.15"
regex==2025.10.23 ; python_version >= "3.10" and python_version < "3.15"
requests==2.32.5 ; python_version >= "3.10" and python_version < "3.15"
safetensors==0.4.5 ; python_version >= "3.10" and python_version < "3.15"
setuptools==80.9.0 ; python_version >= "3.12" and python_version < "3.15"
sympy==1.14.0 ; python_version >= "3.10" and python_version < "3.15"
tokenizers==0.22.1 ; python_version >= "3.10" and python_version < "3.15"
torch==2.9.0 ; python_version >= "3.10" and python_version < "3.15"
torchaudio==2.9.0 ; python_version >= "3.10" and python_version < "3.15"
torchvision==0.24.0 ; python_version >= "3.10" and python_version < "3.15"
tqdm==4.67.1 ; python_version >= "3.10" and python_version < "3.15"
transformers==4.57.1 ; python_version >= "3.10" and python_version < "3.15"
triton==3.5.0 ; python_version >= "3.10" and python_version < "3.15" and platform_system == "Linux" and platform_machine == "x86_64"
typing-extensions==4.15.0 ; python_version >= "3.10" and python_version < "3.15"
urllib3==2.5.0 ; python_version >= "3.10" and python_version < "3.15"
zipp==3.23.0 ; python_version >= "3.10" and python_version < "3.15"
+56
View File
@@ -0,0 +1,56 @@
# Ruff configuration
line-length = 100
target-version = "py310"
src = ["src"]
[lint]
select = [
"E", # pycodestyle errors
"W", # pycodestyle warnings
"F", # pyflakes
"I", # isort
"B", # flake8-bugbear
"C4", # flake8-comprehensions
"UP", # pyupgrade
"D", # pydocstyle
"N", # pep8-naming
"S", # flake8-bandit
"T20", # flake8-print
"PT", # flake8-pytest-style
"RET", # flake8-return
"SIM", # flake8-simplify
"ARG", # flake8-unused-arguments
"PIE", # flake8-pie
"PL", # pylint
]
ignore = [
"E501", # line too long (handled by formatter)
"B008", # do not perform function calls in argument defaults
"C901", # too complex
"PLR0913", # too many arguments
"PLR0912", # too many branches
"D203", # one-blank-line-before-class (conflicts with D211)
"D213", # multi-line-summary-second-line (conflicts with D212)
]
[lint.per-file-ignores]
"__init__.py" = ["F401", "D104", "N999"] # Root __init__.py for ComfyUI-Diffusers
"tests/__init__.py" = ["D104", "N999"] # tests/__init__.py
"src/**/__init__.py" = ["F401", "D104"] # src package __init__.py files
"tests/**/*.py" = ["S101", "D103", "PLR2004", "S106", "ARG001", "PLR0913"]
"src/nodes/*.py" = ["N802", "D102"] # Allow INPUT_TYPES uppercase & missing docstrings
"src/nodes/dto/*.py" = ["N815"] # Allow ClassVar uppercase (COMFY_TYPE)
"src/infra/diffusers/_autoencoder_repository.py" = ["PLR0915"] # Allow large conversion function
"src/di/*.py" = ["D102"] # DI container doesn't need method docstrings
[lint.isort]
known-first-party = ["src"]
section-order = ["future", "standard-library", "third-party", "first-party", "local-folder"]
[lint.pydocstyle]
convention = "google"
[format]
quote-style = "double"
indent-style = "space"
+5
View File
@@ -0,0 +1,5 @@
from ._cache import get_cache_dir
from ._env import load_envs
from ._path import get_project_root
__all__ = ["load_envs", "get_cache_dir", "get_project_root"]
+8
View File
@@ -0,0 +1,8 @@
from ._env import get_str
def get_cache_dir() -> str:
cache_dir = get_str("COMFYUI_DIFFUSERS_CACHE_DIR", "./cache")
if cache_dir is None:
raise ValueError("Cache directory environment variable is not set.")
return cache_dir
+62
View File
@@ -0,0 +1,62 @@
from __future__ import annotations
import os
from functools import lru_cache
from typing import Final
from dotenv import find_dotenv, load_dotenv
from ._path import get_project_root
_TRUE_SET: Final[set[str]] = {"1", "true", "yes", "on"}
_FALSE_SET: Final[set[str]] = {"0", "false", "no", "off"}
@lru_cache(maxsize=1)
def _load_envs_once() -> None:
root = get_project_root()
dotenv_path = str(root / ".env")
default_dotenv_path = str(root / ".env.default")
load_dotenv(find_dotenv(dotenv_path, usecwd=True), override=False)
load_dotenv(find_dotenv(default_dotenv_path, usecwd=True), override=False)
def load_envs(force: bool = False) -> None:
if force:
_load_envs_once.cache_clear()
_load_envs_once()
def get_str(name: str, default: str | None = None) -> str | None:
return os.getenv(name, default)
def get_bool(name: str, default: bool = False) -> bool:
raw = os.getenv(name)
if raw is None:
return default
value = raw.strip().lower()
if value in _TRUE_SET:
return True
if value in _FALSE_SET:
return False
raise ValueError(f"Environment variable '{name}' has invalid boolean value '{raw}'")
def get_int(name: str, default: int | None = None) -> int | None:
raw = os.getenv(name)
if raw is None:
return default
try:
return int(raw.strip())
except (TypeError, ValueError) as err:
raise ValueError(
f"Environment variable '{name}' has invalid integer value '{raw}'"
) from err
def require(name: str) -> str:
value = os.getenv(name)
if value is None:
raise KeyError(f"Required environment variable not set: '{name}'")
return value
+5
View File
@@ -0,0 +1,5 @@
from pathlib import Path
def get_project_root() -> Path:
return Path(__file__).parents[2].resolve()
+3
View File
@@ -0,0 +1,3 @@
from .di import Container
__all__ = ["Container"]
+51
View File
@@ -0,0 +1,51 @@
"""Dependency injection container configuration."""
from dependency_injector import containers, providers
from ..infra.diffusers import (
DiffusersAutoencoderRepository,
DiffusersPipelineRepository,
DiffusersSamplerRepository,
DiffusersSchedulerRepository,
DiffusersTextEncoderRepository,
)
from ..usecase import (
AutoencoderUsecase,
ClipTextEncodeUsecase,
PipelineUsecase,
SamplerUsecase,
SchedulerUsecase,
)
class Container(containers.DeclarativeContainer):
"""Dependency injection container for the application."""
# 1. Repositories
pipeline_repository = providers.Factory(DiffusersPipelineRepository)
autoencoder_repository = providers.Factory(DiffusersAutoencoderRepository)
text_encoder_repository = providers.Factory(DiffusersTextEncoderRepository)
sampler_repository = providers.Factory(DiffusersSamplerRepository)
scheduler_repository = providers.Factory(DiffusersSchedulerRepository)
# 2. Usecases
pipeline_usecase = providers.Factory(
PipelineUsecase,
pipeline_repo=pipeline_repository,
)
autoencoder_usecase = providers.Factory(
AutoencoderUsecase,
autoencoder_repo=autoencoder_repository,
)
clip_text_encode_usecase = providers.Factory(
ClipTextEncodeUsecase,
text_encoder_repo=text_encoder_repository,
)
sampler_usecase = providers.Factory(
SamplerUsecase,
sampler_repo=sampler_repository,
)
scheduler_usecase = providers.Factory(
SchedulerUsecase,
scheduler_repo=scheduler_repository,
)
View File
+23
View File
@@ -0,0 +1,23 @@
from ._autoencoder import Autoencoder
from ._cfg_scale import CFGScale
from ._clip import Clip
from ._conditioning import Conditioning
from ._image import Image
from ._image_size import ImageSize
from ._pipeline import Pipeline
from ._scheduler import Scheduler
from ._seed import Seed
from ._steps import Steps
__all__ = [
"Pipeline",
"Autoencoder",
"Clip",
"Conditioning",
"Image",
"Scheduler",
"ImageSize",
"Steps",
"CFGScale",
"Seed",
]
+19
View File
@@ -0,0 +1,19 @@
from dataclasses import dataclass
from diffusers import AutoencoderKL
@dataclass(frozen=True)
class Autoencoder:
autoencoder: AutoencoderKL
path: str
def __post_init__(self) -> None:
if not isinstance(self.autoencoder, AutoencoderKL):
raise ValueError(
f"Autoencoder must be an instance of AutoencoderKL, got {type(self.autoencoder).__name__}"
)
if not isinstance(self.path, str):
raise ValueError(f"Path must be a string, got {type(self.path).__name__}")
if not self.path.strip():
raise ValueError("Path cannot be empty or whitespace only")
+12
View File
@@ -0,0 +1,12 @@
from dataclasses import dataclass
@dataclass(frozen=True)
class CFGScale:
value: float
def __post_init__(self) -> None:
if isinstance(self.value, bool) or not isinstance(self.value, float):
raise ValueError(f"CFG scale must be a float, got {type(self.value).__name__}")
if self.value < 0.0:
raise ValueError(f"CFG scale must be non-negative, got {self.value}")
+24
View File
@@ -0,0 +1,24 @@
from dataclasses import dataclass
from transformers import CLIPTextModel, CLIPTokenizer
@dataclass(frozen=True)
class Clip:
tokenizer: CLIPTokenizer
text_encoder: CLIPTextModel
path: str
def __post_init__(self) -> None:
if not isinstance(self.tokenizer, CLIPTokenizer):
raise ValueError(
f"Tokenizer must be an instance of CLIPTokenizer, got {type(self.tokenizer).__name__}"
)
if not isinstance(self.text_encoder, CLIPTextModel):
raise ValueError(
f"Text encoder must be an instance of CLIPTextModel, got {type(self.text_encoder).__name__}"
)
if not isinstance(self.path, str):
raise ValueError(f"Path must be a string, got {type(self.path).__name__}")
if not self.path.strip():
raise ValueError("Path cannot be empty or whitespace only")
+28
View File
@@ -0,0 +1,28 @@
from dataclasses import dataclass
import torch
@dataclass(frozen=True)
class Conditioning:
conditioning: torch.Tensor
path: str
prompt: str
def __post_init__(self) -> None:
if not isinstance(self.conditioning, torch.Tensor):
raise ValueError(
f"Conditioning must be a torch.Tensor, got {type(self.conditioning).__name__}"
)
if not isinstance(self.path, str):
raise ValueError(f"Path must be a string, got {type(self.path).__name__}")
if not isinstance(self.prompt, str):
raise ValueError(f"Prompt must be a string, got {type(self.prompt).__name__}")
if torch.isnan(self.conditioning).any():
raise ValueError("Conditioning tensor contains NaN values")
if torch.isinf(self.conditioning).any():
raise ValueError("Conditioning tensor contains Inf values")
if not self.path.strip():
raise ValueError("Path cannot be empty or whitespace only")
if not self.prompt.strip():
raise ValueError("Prompt cannot be empty or whitespace only")
+12
View File
@@ -0,0 +1,12 @@
from dataclasses import dataclass
from PIL import Image as PilImage
@dataclass(frozen=True)
class Image:
image: PilImage.Image
def __post_init__(self) -> None:
if not isinstance(self.image, PilImage.Image):
raise ValueError(f"Image must be a PIL.Image.Image, got {type(self.image).__name__}")
+21
View File
@@ -0,0 +1,21 @@
from dataclasses import dataclass
@dataclass(frozen=True)
class ImageSize:
width: int
height: int
def __post_init__(self) -> None:
if isinstance(self.width, bool) or not isinstance(self.width, int):
raise ValueError(f"Width must be an int, got {type(self.width).__name__}")
if isinstance(self.height, bool) or not isinstance(self.height, int):
raise ValueError(f"Height must be an int, got {type(self.height).__name__}")
if self.width <= 0:
raise ValueError(f"Width must be positive, got {self.width}")
if self.height <= 0:
raise ValueError(f"Height must be positive, got {self.height}")
if self.width % 8 != 0:
raise ValueError(f"Width must be divisible by 8, got {self.width}")
if self.height % 8 != 0:
raise ValueError(f"Height must be divisible by 8, got {self.height}")
+19
View File
@@ -0,0 +1,19 @@
from dataclasses import dataclass
from diffusers import StableDiffusionPipeline
@dataclass(frozen=True)
class Pipeline:
pipeline: StableDiffusionPipeline
path: str
def __post_init__(self) -> None:
if not isinstance(self.pipeline, StableDiffusionPipeline):
raise ValueError(
f"Pipeline must be an instance of StableDiffusionPipeline, got {type(self.pipeline).__name__}"
)
if not isinstance(self.path, str):
raise ValueError(f"Path must be a string, got {type(self.path).__name__}")
if not self.path.strip():
raise ValueError("Path cannot be empty or whitespace only")
+46
View File
@@ -0,0 +1,46 @@
from dataclasses import dataclass
from enum import Enum
from diffusers import (
DDIMScheduler,
DDPMScheduler,
DEISMultistepScheduler,
DPMSolverMultistepScheduler,
DPMSolverSinglestepScheduler,
EulerAncestralDiscreteScheduler,
EulerDiscreteScheduler,
HeunDiscreteScheduler,
KDPM2AncestralDiscreteScheduler,
KDPM2DiscreteScheduler,
UniPCMultistepScheduler,
)
from diffusers.schedulers.scheduling_utils import SchedulerMixin
@dataclass(frozen=True)
class Scheduler:
class Type(Enum):
DDIM = DDIMScheduler
DDPM = DDPMScheduler
DEISMultistep = DEISMultistepScheduler
DPMSolverMultistep = DPMSolverMultistepScheduler
DPMSolverSinglestep = DPMSolverSinglestepScheduler
EulerAncestralDiscrete = EulerAncestralDiscreteScheduler
EulerDiscrete = EulerDiscreteScheduler
HeunDiscrete = HeunDiscreteScheduler
KDPM2AncestralDiscrete = KDPM2AncestralDiscreteScheduler
KDPM2Discrete = KDPM2DiscreteScheduler
UniPCMultistep = UniPCMultistepScheduler
scheduler: SchedulerMixin
path: str
def __post_init__(self) -> None:
if not isinstance(self.scheduler, SchedulerMixin):
raise ValueError(
f"Scheduler must inherit from SchedulerMixin, got {type(self.scheduler).__name__}"
)
if not isinstance(self.path, str):
raise ValueError(f"Path must be a string, got {type(self.path).__name__}")
if not self.path.strip():
raise ValueError("Path cannot be empty or whitespace only")
+10
View File
@@ -0,0 +1,10 @@
from dataclasses import dataclass
@dataclass(frozen=True)
class Seed:
value: int
def __post_init__(self) -> None:
if isinstance(self.value, bool) or not isinstance(self.value, int):
raise ValueError(f"Seed must be an int, got {type(self.value).__name__}")
+12
View File
@@ -0,0 +1,12 @@
from dataclasses import dataclass
@dataclass(frozen=True)
class Steps:
value: int
def __post_init__(self) -> None:
if isinstance(self.value, bool) or not isinstance(self.value, int):
raise ValueError(f"Steps must be an int, got {type(self.value).__name__}")
if self.value <= 0:
raise ValueError(f"Steps must be positive, got {self.value}")
+13
View File
@@ -0,0 +1,13 @@
from ._autoencoder_repository import AutoencoderRepository
from ._pipeline_repository import PipelineRepository
from ._sampler_repository import SamplerRepository
from ._scheduler_repository import SchedulerRepository
from ._text_encoder_repository import TextEncoderRepository
__all__ = [
"PipelineRepository",
"AutoencoderRepository",
"TextEncoderRepository",
"SamplerRepository",
"SchedulerRepository",
]
@@ -0,0 +1,14 @@
from abc import ABC, abstractmethod
import torch
from diffusers import AutoencoderKL
class AutoencoderRepository(ABC):
@abstractmethod
def convert_and_save_from_single_file(self, checkpoint_path: str, dtype: torch.dtype) -> str:
pass
@abstractmethod
def load_autoencoder_from_path(self, model_path: str, dtype: torch.dtype) -> AutoencoderKL:
pass
@@ -0,0 +1,16 @@
from abc import ABC, abstractmethod
import torch
from diffusers import StableDiffusionPipeline
class PipelineRepository(ABC):
@abstractmethod
def convert_and_save_from_single_file(self, checkpoint_path: str, dtype: torch.dtype) -> str:
pass
@abstractmethod
def load_pipeline_from_path(
self, model_path: str, dtype: torch.dtype
) -> StableDiffusionPipeline:
pass
@@ -0,0 +1,25 @@
from abc import ABC, abstractmethod
import torch
from diffusers import AutoencoderKL, StableDiffusionPipeline
from diffusers.schedulers.scheduling_utils import SchedulerMixin
from PIL import Image
from ..model import CFGScale, ImageSize, Seed, Steps
class SamplerRepository(ABC):
@abstractmethod
def sample(
self,
pipeline: StableDiffusionPipeline,
vae: AutoencoderKL,
scheduler: SchedulerMixin,
positive_embeds: torch.Tensor,
negative_embeds: torch.Tensor,
image_size: ImageSize,
steps: Steps,
cfg: CFGScale,
seed: Seed,
) -> list[Image.Image]:
pass
@@ -0,0 +1,14 @@
from abc import ABC, abstractmethod
import torch
from diffusers.schedulers.scheduling_utils import SchedulerMixin
from ...domain.model import Scheduler
class SchedulerRepository(ABC):
@abstractmethod
def create_scheduler(
self, model_path: str, dtype: torch.dtype, scheduler_type: Scheduler.Type
) -> SchedulerMixin:
pass
@@ -0,0 +1,14 @@
from __future__ import annotations
from abc import ABC, abstractmethod
import torch
from transformers import CLIPTextModel, CLIPTokenizer
class TextEncoderRepository(ABC):
@abstractmethod
def encode(
self, tokenizer: CLIPTokenizer, text_encoder: CLIPTextModel, text: str
) -> torch.Tensor:
pass
View File
+13
View File
@@ -0,0 +1,13 @@
from ._autoencoder_repository import DiffusersAutoencoderRepository
from ._pipeline_repository import DiffusersPipelineRepository
from ._sampler_repository import DiffusersSamplerRepository
from ._scheduler_repository import DiffusersSchedulerRepository
from ._text_encoder_repository import DiffusersTextEncoderRepository
__all__ = [
"DiffusersPipelineRepository",
"DiffusersAutoencoderRepository",
"DiffusersTextEncoderRepository",
"DiffusersSamplerRepository",
"DiffusersSchedulerRepository",
]
@@ -0,0 +1,41 @@
import os
import torch
from comfy.model_management import get_torch_device # pyright: ignore[reportMissingImports]
from diffusers import AutoencoderKL
from ...domain.repositories import AutoencoderRepository
from ._cache import get_cache_path, is_vae_cached
class DiffusersAutoencoderRepository(AutoencoderRepository):
def __init__(self) -> None:
self.cache_dir = get_cache_path()
self.device = get_torch_device()
def convert_and_save_from_single_file(self, checkpoint_path: str, dtype: torch.dtype) -> str:
checkpoint_name = os.path.basename(checkpoint_path)
ckpt_cache_path = os.path.join(self.cache_dir, checkpoint_name)
# Check if already cached
if is_vae_cached(ckpt_cache_path):
return ckpt_cache_path
AutoencoderKL.from_single_file(
pretrained_model_link_or_path=checkpoint_path,
torch_dtype=dtype,
cache_dir=self.cache_dir,
).save_pretrained(ckpt_cache_path, safe_serialization=True)
return ckpt_cache_path
def load_autoencoder_from_path(self, model_path: str, dtype: torch.dtype) -> AutoencoderKL:
if not os.path.exists(model_path):
raise FileNotFoundError(f"Model path not found: {model_path}")
vae: AutoencoderKL = AutoencoderKL.from_pretrained( # type: ignore[no-untyped-call]
pretrained_model_name_or_path=model_path,
torch_dtype=dtype,
cache_dir=self.cache_dir,
).to(self.device)
if vae is None:
raise RuntimeError(f"Failed to load AutoencoderKL from path: {model_path}")
return vae
+58
View File
@@ -0,0 +1,58 @@
import os
from pathlib import Path
from typing import Final
from ...config import get_cache_dir, get_project_root
# Centralized cache helpers for diffusers-related repositories
_PIPELINE_INDEX_FILE: Final[str] = "model_index.json"
_VAE_CONFIG_FILE: Final[str] = "config.json"
_VAE_WEIGHT_CANDIDATES: Final[tuple[str, ...]] = (
"diffusion_pytorch_model.safetensors",
"diffusion_pytorch_model.bin",
"pytorch_model.bin",
"model.safetensors",
)
def setup_cache_path() -> Path:
env_cache_str = get_cache_dir()
env_cache_path = Path(env_cache_str)
if env_cache_path.is_absolute():
cache_dir = env_cache_path
else:
project_root = get_project_root()
cache_dir = (project_root / env_cache_path).resolve()
if not cache_dir.exists():
cache_dir.mkdir(parents=True, exist_ok=True)
return cache_dir
def get_cache_path(filename: str | None = None) -> Path:
cache_dir = setup_cache_path()
if filename is None:
return cache_dir
return cache_dir / filename
def is_pipeline_cached(cache_path: str) -> bool:
"""Returns True if the Stable Diffusion pipeline has already been converted.
Criteria: directory exists and contains a `model_index.json`.
"""
return os.path.isdir(cache_path) and os.path.isfile(
os.path.join(cache_path, _PIPELINE_INDEX_FILE)
)
def is_vae_cached(cache_path: str) -> bool:
"""Returns True if the VAE has already been converted.
Criteria: directory exists, contains `config.json`, and at least one known weight file.
"""
if not (
os.path.isdir(cache_path) and os.path.isfile(os.path.join(cache_path, _VAE_CONFIG_FILE))
):
return False
return any(os.path.isfile(os.path.join(cache_path, fname)) for fname in _VAE_WEIGHT_CANDIDATES)
@@ -0,0 +1,39 @@
import os
import torch
from comfy.model_management import get_torch_device # pyright: ignore[reportMissingImports]
from diffusers import StableDiffusionPipeline
from ...domain.repositories import PipelineRepository
from ._cache import get_cache_path, is_pipeline_cached
class DiffusersPipelineRepository(PipelineRepository):
def __init__(self) -> None:
self.cache_dir = get_cache_path()
self.device = get_torch_device()
def convert_and_save_from_single_file(self, checkpoint_path: str, dtype: torch.dtype) -> str:
checkpoint_name = os.path.basename(checkpoint_path)
ckpt_cache_path = os.path.join(self.cache_dir, checkpoint_name)
# Check if already cached
if is_pipeline_cached(ckpt_cache_path):
return ckpt_cache_path
StableDiffusionPipeline.from_single_file(
pretrained_model_link_or_path=checkpoint_path,
torch_dtype=dtype,
cache_dir=self.cache_dir,
).save_pretrained(ckpt_cache_path, safe_serialization=True)
return ckpt_cache_path
def load_pipeline_from_path(
self, model_path: str, dtype: torch.dtype
) -> StableDiffusionPipeline:
pipe: StableDiffusionPipeline = StableDiffusionPipeline.from_pretrained( # type: ignore[no-untyped-call]
pretrained_model_name_or_path=model_path,
torch_dtype=dtype,
cache_dir=self.cache_dir,
).to(self.device)
return pipe
@@ -0,0 +1,39 @@
import torch
from comfy.model_management import get_torch_device # pyright: ignore[reportMissingImports]
from diffusers import AutoencoderKL, StableDiffusionPipeline
from diffusers.schedulers.scheduling_utils import SchedulerMixin
from PIL import Image
from ...domain.model import CFGScale, ImageSize, Seed, Steps
from ...domain.repositories import SamplerRepository
class DiffusersSamplerRepository(SamplerRepository):
def __init__(self) -> None:
self.device = get_torch_device()
def sample(
self,
pipeline: StableDiffusionPipeline,
vae: AutoencoderKL,
scheduler: SchedulerMixin,
positive_embeds: torch.Tensor,
negative_embeds: torch.Tensor,
image_size: ImageSize,
steps: Steps,
cfg: CFGScale,
seed: Seed,
) -> list[Image.Image]:
result = pipeline( # type: ignore[operator]
prompt_embeds=positive_embeds,
vae=vae,
scheduler=scheduler,
height=image_size.height,
width=image_size.width,
num_inference_steps=steps.value,
guidance_scale=cfg.value,
negative_prompt_embeds=negative_embeds,
generator=torch.Generator(self.device).manual_seed(seed.value),
)
images: list[Image.Image] = result.images
return images
@@ -0,0 +1,22 @@
import torch
from diffusers.schedulers.scheduling_utils import SchedulerMixin
from ...domain.model import Scheduler
from ...domain.repositories import SchedulerRepository
from ._cache import get_cache_path
class DiffusersSchedulerRepository(SchedulerRepository):
def __init__(self) -> None:
self.cache_dir = get_cache_path()
def create_scheduler(
self, model_path: str, dtype: torch.dtype, scheduler_type: Scheduler.Type
) -> SchedulerMixin:
scheduler: SchedulerMixin = scheduler_type.value.from_pretrained(
pretrained_model_name_or_path=model_path,
subfolder="scheduler",
torch_dtype=dtype,
cache_dir=self.cache_dir,
)
return scheduler
@@ -0,0 +1,46 @@
from __future__ import annotations
from typing import Any
import torch
from comfy.model_management import get_torch_device # pyright: ignore[reportMissingImports]
from transformers import CLIPTextModel, CLIPTokenizer
from ...domain.repositories._text_encoder_repository import TextEncoderRepository
class DiffusersTextEncoderRepository(TextEncoderRepository):
def __init__(self) -> None:
self.device = get_torch_device()
def encode(
self, tokenizer: CLIPTokenizer, text_encoder: CLIPTextModel, text: str
) -> torch.Tensor:
max_length = tokenizer.model_max_length
input_ids = tokenizer(text, return_tensors="pt").input_ids
text_length = input_ids.shape[-1]
if max_length < text_length:
text_ids = input_ids.to(self.device)
else:
text_ids = tokenizer(
text,
truncation=False,
padding="max_length",
max_length=max_length,
return_tensors="pt",
).input_ids.to(self.device)
concat_embeds: list[torch.Tensor] = []
with torch.no_grad():
for i in range(0, text_ids.shape[-1], max_length):
segment_ids = text_ids[:, i : i + max_length]
outputs: Any = text_encoder(segment_ids)
if hasattr(outputs, "last_hidden_state"):
embeds = outputs.last_hidden_state
else:
embeds = outputs[0]
concat_embeds.append(embeds)
text_embeds: torch.Tensor = torch.cat(concat_embeds, dim=1)
return text_embeds
+41
View File
@@ -0,0 +1,41 @@
from ..di import Container
from ._diffusers_clip_text_encode import DiffusersClipTextEncode
from ._diffusers_pipeline_loader import DiffusersPipelineLoader
from ._diffusers_sampler import DiffusersSampler
from ._diffusers_scheduler_loader import DiffusersSchedulerLoader
from ._diffusers_vae_loader import DiffusersVaeLoader
from .dto import (
ComfyUIAutoencoderDTO,
ComfyUIClipDTO,
ComfyUIConditioningDTO,
ComfyUIPipelineDTO,
ComfyUISchedulerDTO,
)
container = Container()
container.wire(modules=[__name__])
NODE_CLASS_MAPPINGS = {
DiffusersPipelineLoader.__name__: DiffusersPipelineLoader,
DiffusersVaeLoader.__name__: DiffusersVaeLoader,
DiffusersClipTextEncode.__name__: DiffusersClipTextEncode,
DiffusersSampler.__name__: DiffusersSampler,
DiffusersSchedulerLoader.__name__: DiffusersSchedulerLoader,
}
NODE_DISPLAY_NAME_MAPPINGS = {
DiffusersPipelineLoader.__name__: "Diffusers Pipeline Loader",
DiffusersVaeLoader.__name__: "Diffusers VAE Loader",
DiffusersClipTextEncode.__name__: "Diffusers CLIP Text Encode",
DiffusersSampler.__name__: "Diffusers Sampler",
DiffusersSchedulerLoader.__name__: "Diffusers Scheduler Loader",
}
__all__ = [
"NODE_CLASS_MAPPINGS",
"NODE_DISPLAY_NAME_MAPPINGS",
"ComfyUIPipelineDTO",
"ComfyUIAutoencoderDTO",
"ComfyUIClipDTO",
"ComfyUIConditioningDTO",
"ComfyUISchedulerDTO",
]
+40
View File
@@ -0,0 +1,40 @@
from dependency_injector.wiring import Provide, inject
from ..di import Container
from ..usecase import ClipTextEncodeUsecase
from .dto import ComfyUIClipDTO, ComfyUIConditioningDTO
class DiffusersClipTextEncode:
def __init__(self) -> None:
pass
@classmethod
def INPUT_TYPES(cls) -> dict[str, dict[str, tuple[str, ...] | tuple[str, dict[str, bool]]]]:
return {
"required": {
"clip": (ComfyUIClipDTO.COMFY_TYPE,),
"text": ("STRING", {"multiline": True}),
},
}
RETURN_TYPES = (ComfyUIConditioningDTO.COMFY_TYPE,)
FUNCTION = "execute"
CATEGORY = "Diffusers"
@inject
def execute(
self,
clip: ComfyUIClipDTO,
text: str,
usecase: ClipTextEncodeUsecase = Provide[Container.clip_text_encode_usecase],
) -> tuple[ComfyUIConditioningDTO]:
clip_domain = ComfyUIClipDTO.to_domain(clip)
conditioning_model = usecase.execute(
clip_domain.tokenizer,
clip_domain.text_encoder,
clip_domain.path,
text,
)
conditioning_dto = ComfyUIConditioningDTO.from_domain(conditioning_model)
return (conditioning_dto,)
+41
View File
@@ -0,0 +1,41 @@
import folder_paths # pyright: ignore[reportMissingImports]
from dependency_injector.wiring import Provide, inject
from ..di import Container
from ..usecase import PipelineUsecase
from .dto import ComfyUIClipDTO, ComfyUIPipelineDTO
class DiffusersPipelineLoader:
def __init__(self) -> None:
pass
@classmethod
def INPUT_TYPES(cls) -> dict[str, dict[str, tuple[tuple[str, ...], ...]]]:
return {
"required": {
"checkpoint_name": (folder_paths.get_filename_list("checkpoints"),),
}
}
RETURN_TYPES = (
ComfyUIPipelineDTO.COMFY_TYPE,
ComfyUIClipDTO.COMFY_TYPE,
)
FUNCTION = "execute"
CATEGORY = "Diffusers"
@inject
def execute(
self,
checkpoint_name: str,
usecase: PipelineUsecase = Provide[Container.pipeline_usecase],
) -> tuple[ComfyUIPipelineDTO, ComfyUIClipDTO]:
checkpoint_path = folder_paths.get_full_path("checkpoints", checkpoint_name)
pipeline_model, clip_model = usecase.execute(checkpoint_path)
pipeline_dto = ComfyUIPipelineDTO.from_domain(pipeline_model)
clip_dto = ComfyUIClipDTO.from_domain(clip_model)
return (
pipeline_dto,
clip_dto,
)
+84
View File
@@ -0,0 +1,84 @@
from dependency_injector.wiring import Provide, inject
from ..di import Container
from ..domain.model import CFGScale, ImageSize, Seed, Steps
from ..usecase import SamplerUsecase
from .dto import (
ComfyUIAutoencoderDTO,
ComfyUIConditioningDTO,
ComfyUIImage,
ComfyUIImageDTO,
ComfyUIPipelineDTO,
ComfyUISchedulerDTO,
)
class DiffusersSampler:
def __init__(self) -> None:
pass
@classmethod
def INPUT_TYPES(
cls,
) -> dict[str, dict[str, tuple[str, ...] | tuple[str, dict[str, int | float]]]]:
return {
"required": {
"pipeline": (ComfyUIPipelineDTO.COMFY_TYPE,),
"vae": (ComfyUIAutoencoderDTO.COMFY_TYPE,),
"scheduler": (ComfyUISchedulerDTO.COMFY_TYPE,),
"positive_embeds": (ComfyUIConditioningDTO.COMFY_TYPE,),
"negative_embeds": (ComfyUIConditioningDTO.COMFY_TYPE,),
"width": ("INT", {"default": 512, "min": 8, "max": 8192, "step": 8.0}),
"height": ("INT", {"default": 512, "min": 8, "max": 8192, "step": 8.0}),
"steps": ("INT", {"default": 20, "min": 1, "max": 10000, "step": 1.0}),
"cfg": (
"FLOAT",
{"default": 8.0, "min": 0.0, "max": 100.0, "step": 0.1, "round": 0.01},
),
"seed": ("INT", {"default": 0, "min": 0, "max": 0xFFFFFFFFFFFFFFFF, "step": 1.0}),
},
}
RETURN_TYPES = (ComfyUIImage.COMFY_TYPE,)
FUNCTION = "execute"
CATEGORY = "Diffusers"
@inject
def execute(
self,
pipeline: ComfyUIPipelineDTO,
vae: ComfyUIAutoencoderDTO,
scheduler: ComfyUISchedulerDTO,
positive_embeds: ComfyUIConditioningDTO,
negative_embeds: ComfyUIConditioningDTO,
width: int,
height: int,
steps: int,
cfg: float,
seed: int,
usecase: SamplerUsecase = Provide[Container.sampler_usecase],
) -> tuple[ComfyUIImageDTO]:
pipeline_domain = ComfyUIPipelineDTO.to_domain(pipeline)
vae_domain = ComfyUIAutoencoderDTO.to_domain(vae)
scheduler_domain = ComfyUISchedulerDTO.to_domain(scheduler)
positive_embeds_domain = ComfyUIConditioningDTO.to_domain(positive_embeds)
negative_embeds_domain = ComfyUIConditioningDTO.to_domain(negative_embeds)
image_size = ImageSize(width=width, height=height)
steps_vo = Steps(value=steps)
cfg_vo = CFGScale(value=cfg)
seed_vo = Seed(value=seed)
images_model = usecase.execute(
pipeline_domain,
vae_domain,
scheduler_domain,
positive_embeds_domain,
negative_embeds_domain,
image_size,
steps_vo,
cfg_vo,
seed_vo,
)
sampler_dto = ComfyUIImage.from_domains(images_model)
return (sampler_dto,)
+37
View File
@@ -0,0 +1,37 @@
from dependency_injector.wiring import Provide, inject
from ..di import Container
from ..usecase import SchedulerUsecase
from .dto import ComfyUIPipelineDTO, ComfyUISchedulerDTO
from .map import SchedulerMap
class DiffusersSchedulerLoader:
def __init__(self) -> None:
pass
@classmethod
def INPUT_TYPES(cls) -> dict[str, dict[str, tuple[str] | tuple[list[str]]]]:
return {
"required": {
"pipeline": (ComfyUIPipelineDTO.COMFY_TYPE,),
"scheduler_name": (SchedulerMap.SCHEDULERS,),
}
}
RETURN_TYPES = (ComfyUISchedulerDTO.COMFY_TYPE,)
FUNCTION = "execute"
CATEGORY = "Diffusers"
@inject
def execute(
self,
pipeline: ComfyUIPipelineDTO,
scheduler_name: str,
usecase: SchedulerUsecase = Provide[Container.scheduler_usecase],
) -> tuple[ComfyUISchedulerDTO]:
domain = ComfyUIPipelineDTO.to_domain(pipeline)
scheduler_type = SchedulerMap.to_domain(scheduler_name)
scheduler_domain = usecase.execute(domain, scheduler_type)
scheduler_dto = ComfyUISchedulerDTO.from_domain(scheduler_domain)
return (scheduler_dto,)
+34
View File
@@ -0,0 +1,34 @@
import folder_paths # pyright: ignore[reportMissingImports]
from dependency_injector.wiring import Provide, inject
from ..di import Container
from ..usecase import AutoencoderUsecase
from .dto import ComfyUIAutoencoderDTO
class DiffusersVaeLoader:
def __init__(self) -> None:
pass
@classmethod
def INPUT_TYPES(cls) -> dict[str, dict[str, tuple[tuple[str, ...], ...]]]:
return {
"required": {
"vae_name": (folder_paths.get_filename_list("vae"),),
},
}
RETURN_TYPES = (ComfyUIAutoencoderDTO.COMFY_TYPE,)
FUNCTION = "execute"
CATEGORY = "Diffusers"
@inject
def execute(
self,
vae_name: str,
usecase: AutoencoderUsecase = Provide[Container.autoencoder_usecase],
) -> tuple[ComfyUIAutoencoderDTO]:
vae_path = folder_paths.get_full_path("vae", vae_name)
autoencoder_model = usecase.execute(vae_path)
autoencoder_dto = ComfyUIAutoencoderDTO.from_domain(autoencoder_model)
return (autoencoder_dto,)
+16
View File
@@ -0,0 +1,16 @@
from ._autoencoder_dto import ComfyUIAutoencoderDTO
from ._clip import ComfyUIClipDTO
from ._comfyui import ComfyUIImage, ComfyUIImageDTO
from ._conditioning_dto import ComfyUIConditioningDTO
from ._pipeline_dto import ComfyUIPipelineDTO
from ._scheduler_dto import ComfyUISchedulerDTO
__all__ = [
"ComfyUIPipelineDTO",
"ComfyUIAutoencoderDTO",
"ComfyUIClipDTO",
"ComfyUIConditioningDTO",
"ComfyUIImage",
"ComfyUIImageDTO",
"ComfyUISchedulerDTO",
]
+69
View File
@@ -0,0 +1,69 @@
from dataclasses import dataclass
from typing import ClassVar
from diffusers import AutoencoderKL
from ...domain.model import Autoencoder
@dataclass
class ComfyUIAutoencoderDTO:
"""Data Transfer Object for Diffusers Autoencoder in ComfyUI.
This DTO is designed to be used across different custom nodes.
Other node developers can import and use this type for autoencoder operations.
Attributes:
autoencoder: AutoencoderKL instance from diffusers library
path: Path to the loaded model directory
Class Attributes:
COMFY_TYPE: Type name used in ComfyUI's RETURN_TYPES and INPUT_TYPES
Example:
```python
from your_node_package.nodes.dto import ComfyUIAutoencoderDTO
class YourCustomNode:
@classmethod
def INPUT_TYPES(cls):
return {"required": {"autoencoder": (ComfyUIAutoencoderDTO.COMFY_TYPE,)}}
RETURN_TYPES = (ComfyUIAutoencoderDTO.COMFY_TYPE,)
def execute(self, autoencoder: ComfyUIAutoencoderDTO):
# Use the autoencoder for encoding/decoding
encoded = autoencoder.autoencoder.encode(image)
return (encoded,)
```
"""
# ClassVar to avoid being treated as a dataclass field
COMFY_TYPE: ClassVar[str] = "DIFFUSERS_VAE"
autoencoder: AutoencoderKL
path: str
@classmethod
def from_domain(cls, autoencoder: Autoencoder) -> "ComfyUIAutoencoderDTO":
"""Create DTO from domain model.
Args:
autoencoder: Autoencoder from domain layer
Returns:
ComfyUIAutoencoderDTO instance
"""
return cls(autoencoder=autoencoder.autoencoder, path=autoencoder.path)
@classmethod
def to_domain(cls, dto: "ComfyUIAutoencoderDTO") -> "Autoencoder":
"""Convert DTO back to domain model.
Args:
dto: ComfyUIAutoencoderDTO instance
Returns:
Autoencoder domain model
"""
return Autoencoder(autoencoder=dto.autoencoder, path=dto.path)
+76
View File
@@ -0,0 +1,76 @@
from dataclasses import dataclass
from typing import ClassVar
from transformers import CLIPTextModel, CLIPTokenizer
from ...domain.model import Clip
@dataclass
class ComfyUIClipDTO:
"""Data Transfer Object for CLIP (Tokenizer and Text Encoder) in ComfyUI.
This DTO is designed to be used across different custom nodes.
Other node developers can import and use this type for CLIP operations.
Attributes:
tokenizer: CLIPTokenizer instance from transformers library
text_encoder: CLIPTextModel instance from transformers library
path: Path to the loaded model directory
Class Attributes:
COMFY_TYPE: Type name used in ComfyUI's RETURN_TYPES and INPUT_TYPES
Example:
```python
from your_node_package.nodes.dto import ComfyUIClipDTO
class YourCustomNode:
@classmethod
def INPUT_TYPES(cls):
return {"required": {"clip": (ComfyUIClipDTO.COMFY_TYPE,)}}
RETURN_TYPES = (ComfyUIClipDTO.COMFY_TYPE,)
def execute(self, clip: ComfyUIClipDTO):
# Use the tokenizer and text encoder
tokens = clip.tokenizer(prompt, return_tensors="pt")
embeddings = clip.text_encoder(**tokens)
return (embeddings,)
```
"""
# ClassVar to avoid being treated as a dataclass field
COMFY_TYPE: ClassVar[str] = "DIFFUSERS_CLIP"
tokenizer: CLIPTokenizer
text_encoder: CLIPTextModel
path: str
@classmethod
def from_domain(cls, clip: Clip) -> "ComfyUIClipDTO":
"""Create DTO from domain model.
Args:
clip: Clip from domain layer
Returns:
ComfyUIClipDTO instance
"""
return cls(
tokenizer=clip.tokenizer,
text_encoder=clip.text_encoder,
path=clip.path,
)
@classmethod
def to_domain(cls, dto: "ComfyUIClipDTO") -> "Clip":
"""Convert DTO back to domain model.
Args:
dto: ComfyUIClipDTO instance
Returns:
Clip domain model
"""
return Clip(tokenizer=dto.tokenizer, text_encoder=dto.text_encoder, path=dto.path)
+6
View File
@@ -0,0 +1,6 @@
from ._image_dto import ComfyUIImage, ComfyUIImageDTO
__all__ = [
"ComfyUIImageDTO",
"ComfyUIImage",
]
+23
View File
@@ -0,0 +1,23 @@
from dataclasses import dataclass
from typing import TypeAlias
import torch
from torchvision.transforms import ToTensor
from ....domain.model import Image
ComfyUIImageDTO: TypeAlias = torch.Tensor
@dataclass
class ComfyUIImage:
COMFY_TYPE: str = "IMAGE"
@classmethod
def from_domain(cls, image: Image) -> "ComfyUIImageDTO":
tensor: torch.Tensor = ToTensor()(image.image)
return tensor.permute(1, 2, 0)
@classmethod
def from_domains(cls, images: list[Image]) -> "ComfyUIImageDTO":
return torch.stack([cls.from_domain(image) for image in images])
+75
View File
@@ -0,0 +1,75 @@
from dataclasses import dataclass
from typing import ClassVar
import torch
from ...domain.model import Conditioning
@dataclass
class ComfyUIConditioningDTO:
"""Data Transfer Object for Diffusers Conditioning in ComfyUI.
This DTO is designed to be used across different custom nodes.
Other node developers can import and use this type for conditioning operations.
Attributes:
conditioning: Tensor representing text embeddings/conditioning
path: Path to the source model directory
prompt: Prompt text that produced the conditioning
Class Attributes:
COMFY_TYPE: Type name used in ComfyUI's RETURN_TYPES and INPUT_TYPES
Example:
```python
from your_node_package.nodes.dto import ComfyUIConditioningDTO
class YourCustomNode:
@classmethod
def INPUT_TYPES(cls):
return {"required": {"conditioning": (ComfyUIConditioningDTO.COMFY_TYPE,)}}
RETURN_TYPES = (ComfyUIConditioningDTO.COMFY_TYPE,)
def execute(self, conditioning: ComfyUIConditioningDTO):
# Use the conditioning tensor for generation
result = process_with_conditioning(conditioning.conditioning)
return (result,)
```
"""
# ClassVar to avoid being treated as a dataclass field
COMFY_TYPE: ClassVar[str] = "DIFFUSERS_CONDITIONING"
conditioning: torch.Tensor
path: str
prompt: str
@classmethod
def from_domain(cls, conditioning: Conditioning) -> "ComfyUIConditioningDTO":
"""Create DTO from domain model.
Args:
conditioning: Conditioning from domain layer
Returns:
ComfyUIConditioningDTO instance
"""
return cls(
conditioning=conditioning.conditioning,
path=conditioning.path,
prompt=conditioning.prompt,
)
@classmethod
def to_domain(cls, dto: "ComfyUIConditioningDTO") -> Conditioning:
"""Convert DTO back to domain model.
Args:
dto: ComfyUIConditioningDTO instance
Returns:
Conditioning domain model instance
"""
return Conditioning(conditioning=dto.conditioning, path=dto.path, prompt=dto.prompt)
+68
View File
@@ -0,0 +1,68 @@
from dataclasses import dataclass
from typing import ClassVar
from diffusers import StableDiffusionPipeline
from ...domain.model import Pipeline
@dataclass
class ComfyUIPipelineDTO:
"""Data Transfer Object for Diffusers Pipeline in ComfyUI.
This DTO is designed to be used across different custom nodes.
Other node developers can import and use this type for pipeline operations.
Attributes:
pipeline: StableDiffusionPipeline instance from diffusers library
path: Path to the loaded model directory
Class Attributes:
COMFY_TYPE: Type name used in ComfyUI's RETURN_TYPES and INPUT_TYPES
Example:
```python
from your_node_package.nodes.dto import ComfyUIPipelineDTO
class YourCustomNode:
@classmethod
def INPUT_TYPES(cls):
return {"required": {"pipeline": (ComfyUIPipelineDTO.COMFY_TYPE,)}}
RETURN_TYPES = (ComfyUIPipelineDTO.COMFY_TYPE,)
def execute(self, pipeline: ComfyUIPipelineDTO):
result = pipeline.pipeline(prompt="...")
return (result,)
```
"""
# ClassVar to avoid being treated as a dataclass field
COMFY_TYPE: ClassVar[str] = "DIFFUSERS_PIPELINE"
pipeline: StableDiffusionPipeline
path: str
@classmethod
def from_domain(cls, pipeline: Pipeline) -> "ComfyUIPipelineDTO":
"""Create DTO from domain model.
Args:
pipeline: Pipeline from domain layer
Returns:
ComfyUIPipelineDTO instance
"""
return cls(pipeline=pipeline.pipeline, path=pipeline.path)
@classmethod
def to_domain(cls, dto: "ComfyUIPipelineDTO") -> Pipeline:
"""Convert DTO back to domain model.
Args:
dto: ComfyUIPipelineDTO instance
Returns:
Pipeline domain model instance
"""
return Pipeline(pipeline=dto.pipeline, path=dto.path)
+66
View File
@@ -0,0 +1,66 @@
from dataclasses import dataclass
from typing import ClassVar
from diffusers.schedulers.scheduling_utils import SchedulerMixin
from ...domain.model import Scheduler
@dataclass
class ComfyUISchedulerDTO:
"""Data Transfer Object for Diffusers Scheduler in ComfyUI.
This DTO is designed to be used across different custom nodes.
Other node developers can import and use this type for scheduler operations.
Attributes:
scheduler: SchedulerMixin instance from diffusers library
path: Path to the source model directory
Class Attributes:
COMFY_TYPE: Type name used in ComfyUI's RETURN_TYPES and INPUT_TYPES
Example:
```python
from your_node_package.nodes.dto import ComfyUISchedulerDTO
class YourCustomNode:
@classmethod
def INPUT_TYPES(cls):
return {"required": {"scheduler": (ComfyUISchedulerDTO.COMFY_TYPE,)}}
def execute(self, scheduler: ComfyUISchedulerDTO):
# Use the scheduler for sampler operations
pass
```
"""
# ClassVar to avoid being treated as a dataclass field
COMFY_TYPE: ClassVar[str] = "DIFFUSERS_SCHEDULER"
scheduler: SchedulerMixin
path: str
@classmethod
def from_domain(cls, scheduler: Scheduler) -> "ComfyUISchedulerDTO":
"""Create DTO from domain model.
Args:
scheduler: Scheduler from domain layer
Returns:
ComfyUISchedulerDTO instance
"""
return cls(scheduler=scheduler.scheduler, path=scheduler.path)
@classmethod
def to_domain(cls, dto: "ComfyUISchedulerDTO") -> Scheduler:
"""Convert DTO back to domain model.
Args:
dto: ComfyUISchedulerDTO instance
Returns:
Scheduler domain model
"""
return Scheduler(scheduler=dto.scheduler, path=dto.path)
+5
View File
@@ -0,0 +1,5 @@
from ._scheduler_map import SchedulerMap
__all__ = [
"SchedulerMap",
]
+76
View File
@@ -0,0 +1,76 @@
from typing import Final
from ...domain.model import Scheduler
class SchedulerMap:
"""Maps between domain Scheduler.Type enum and ComfyUI string representation.
This map provides bidirectional conversion between the domain layer's
type-safe enum and ComfyUI's string-based interface requirements.
"""
SCHEDULERS: Final[list[str]] = [
"DDIM",
"DDPM",
"DEISMultistep",
"DPMSolverMultistep",
"DPMSolverSinglestep",
"EulerAncestralDiscrete",
"EulerDiscrete",
"HeunDiscrete",
"KDPM2AncestralDiscrete",
"KDPM2Discrete",
"UniPCMultistep",
]
_DOMAIN_TO_STRING: dict[Scheduler.Type, str] = {
Scheduler.Type.DDIM: "DDIM",
Scheduler.Type.DDPM: "DDPM",
Scheduler.Type.DEISMultistep: "DEISMultistep",
Scheduler.Type.DPMSolverMultistep: "DPMSolverMultistep",
Scheduler.Type.DPMSolverSinglestep: "DPMSolverSinglestep",
Scheduler.Type.EulerAncestralDiscrete: "EulerAncestralDiscrete",
Scheduler.Type.EulerDiscrete: "EulerDiscrete",
Scheduler.Type.HeunDiscrete: "HeunDiscrete",
Scheduler.Type.KDPM2AncestralDiscrete: "KDPM2AncestralDiscrete",
Scheduler.Type.KDPM2Discrete: "KDPM2Discrete",
Scheduler.Type.UniPCMultistep: "UniPCMultistep",
}
_STRING_TO_DOMAIN: dict[str, Scheduler.Type] = {v: k for k, v in _DOMAIN_TO_STRING.items()}
@classmethod
def from_domain(cls, scheduler_type: Scheduler.Type) -> str:
"""Convert domain Scheduler.Type enum to ComfyUI string representation.
Args:
scheduler_type: Scheduler.Type enum value from domain layer
Returns:
String representation for ComfyUI interface
Raises:
KeyError: If the scheduler_type is not in the mapping
"""
return cls._DOMAIN_TO_STRING[scheduler_type]
@classmethod
def to_domain(cls, scheduler_name: str) -> Scheduler.Type:
"""Convert ComfyUI string to domain Scheduler.Type enum.
Args:
scheduler_name: String representation from ComfyUI interface
Returns:
Corresponding Scheduler.Type enum value from domain layer
Raises:
ValueError: If the scheduler_name is not recognized
"""
if scheduler_name not in cls._STRING_TO_DOMAIN:
raise ValueError(
f"Unknown scheduler type: '{scheduler_name}'. "
f"Available types: {', '.join(cls.SCHEDULERS)}"
)
return cls._STRING_TO_DOMAIN[scheduler_name]
+13
View File
@@ -0,0 +1,13 @@
from ._autoencoder_usecase import AutoencoderUsecase
from ._clip_text_encode_usecase import ClipTextEncodeUsecase
from ._pipeline_usecase import PipelineUsecase
from ._sampler_usecase import SamplerUsecase
from ._scheduler_usecase import SchedulerUsecase
__all__ = [
"PipelineUsecase",
"AutoencoderUsecase",
"ClipTextEncodeUsecase",
"SamplerUsecase",
"SchedulerUsecase",
]
+26
View File
@@ -0,0 +1,26 @@
import os
import torch
from ..domain.model import Autoencoder
from ..domain.repositories import AutoencoderRepository
class AutoencoderUsecase:
def __init__(self, autoencoder_repo: AutoencoderRepository) -> None:
self.autoencoder_repo = autoencoder_repo
self.dtype = torch.float32
def execute(self, vae_path: str) -> Autoencoder:
if not os.path.exists(vae_path):
raise FileNotFoundError(f"VAE file not found: {vae_path}")
path = self.autoencoder_repo.convert_and_save_from_single_file(vae_path, self.dtype)
if not os.path.exists(path):
raise RuntimeError(f"Failed to convert VAE checkpoint: {vae_path}")
vae = self.autoencoder_repo.load_autoencoder_from_path(path, self.dtype)
if vae is None:
raise RuntimeError(f"Failed to create autoencoder from checkpoint: {vae_path}")
return Autoencoder(autoencoder=vae, path=path)
+27
View File
@@ -0,0 +1,27 @@
from __future__ import annotations
from transformers import CLIPTextModel, CLIPTokenizer
from ..domain.model import Conditioning
from ..domain.repositories._text_encoder_repository import TextEncoderRepository
class ClipTextEncodeUsecase:
def __init__(self, text_encoder_repo: TextEncoderRepository) -> None:
self.text_encoder_repo = text_encoder_repo
def execute(
self,
tokenizer: CLIPTokenizer,
text_encoder: CLIPTextModel,
path: str,
text: str,
) -> Conditioning:
if not text.strip():
raise ValueError("Text input is empty or whitespace only.")
embeddings = self.text_encoder_repo.encode(tokenizer, text_encoder, text)
if embeddings is None:
raise RuntimeError("Failed to generate embeddings.")
return Conditioning(conditioning=embeddings, path=path, prompt=text)
+33
View File
@@ -0,0 +1,33 @@
import os
import torch
from ..domain.model import Clip, Pipeline
from ..domain.repositories import PipelineRepository
class PipelineUsecase:
def __init__(self, pipeline_repo: PipelineRepository) -> None:
self.pipeline_repo = pipeline_repo
self.dtype = torch.float32
def execute(self, checkpoint_path: str) -> tuple[Pipeline, Clip]:
if not os.path.exists(checkpoint_path):
raise FileNotFoundError(f"Checkpoint file not found: {checkpoint_path}")
path = self.pipeline_repo.convert_and_save_from_single_file(checkpoint_path, self.dtype)
if not os.path.exists(path):
raise RuntimeError(f"Failed to convert pipeline checkpoint: {checkpoint_path}")
pipe = self.pipeline_repo.load_pipeline_from_path(path, self.dtype)
if pipe is None:
raise RuntimeError(f"Failed to create pipeline from checkpoint: {checkpoint_path}")
pipeline = Pipeline(pipeline=pipe, path=path)
clip = Clip(
tokenizer=pipe.tokenizer, # type: ignore[attr-defined]
text_encoder=pipe.text_encoder, # type: ignore[attr-defined]
path=checkpoint_path,
)
return pipeline, clip
+51
View File
@@ -0,0 +1,51 @@
from __future__ import annotations
import torch
from ..domain.model import (
Autoencoder,
CFGScale,
Conditioning,
Image,
ImageSize,
Pipeline,
Scheduler,
Seed,
Steps,
)
from ..domain.repositories import SamplerRepository
class SamplerUsecase:
def __init__(self, sampler_repo: SamplerRepository) -> None:
self.sampler_repo = sampler_repo
self.dtype = torch.float32
def execute(
self,
pipeline: Pipeline,
vae: Autoencoder,
scheduler: Scheduler,
positive_embeds: Conditioning,
negative_embeds: Conditioning,
image_size: ImageSize,
steps: Steps,
cfg: CFGScale,
seed: Seed,
) -> list[Image]:
images = self.sampler_repo.sample(
pipeline.pipeline,
vae.autoencoder,
scheduler.scheduler,
positive_embeds.conditioning,
negative_embeds.conditioning,
image_size,
steps,
cfg,
seed,
)
if not images:
raise RuntimeError("Sampler repository returned no images.")
return [Image(image=img) for img in images]
+20
View File
@@ -0,0 +1,20 @@
import torch
from diffusers.schedulers.scheduling_utils import SchedulerMixin
from ..domain.model import Pipeline, Scheduler
from ..domain.repositories import SchedulerRepository
class SchedulerUsecase:
def __init__(self, scheduler_repo: SchedulerRepository) -> None:
self.scheduler_repo = scheduler_repo
self.dtype = torch.float32
def execute(self, pipeline: Pipeline, scheduler_type: Scheduler.Type) -> Scheduler:
scheduler_obj: SchedulerMixin = self.scheduler_repo.create_scheduler(
pipeline.path,
self.dtype,
scheduler_type,
)
return Scheduler(scheduler=scheduler_obj, path=pipeline.path)
View File
-213
View File
@@ -1,213 +0,0 @@
import io
import torch
import requests
import numpy as np
from PIL import Image
from omegaconf import OmegaConf
from torchvision.transforms import ToTensor
from diffusers.pipelines.stable_diffusion.convert_from_ckpt import (
assign_to_checkpoint,
conv_attn_to_linear,
create_vae_diffusers_config,
renew_vae_attention_paths,
renew_vae_resnet_paths,
)
from diffusers import (
AutoencoderKL,
DDIMScheduler,
DDPMScheduler,
DEISMultistepScheduler,
DPMSolverMultistepScheduler,
DPMSolverSinglestepScheduler,
EulerAncestralDiscreteScheduler,
EulerDiscreteScheduler,
HeunDiscreteScheduler,
KDPM2AncestralDiscreteScheduler,
KDPM2DiscreteScheduler,
UniPCMultistepScheduler,
)
SCHEDULERS = {
'DDIM' : DDIMScheduler,
'DDPM' : DDPMScheduler,
'DEISMultistep' : DEISMultistepScheduler,
'DPMSolverMultistep' : DPMSolverMultistepScheduler,
'DPMSolverSinglestep' : DPMSolverSinglestepScheduler,
'EulerAncestralDiscrete' : EulerAncestralDiscreteScheduler,
'EulerDiscrete' : EulerDiscreteScheduler,
'HeunDiscrete' : HeunDiscreteScheduler,
'KDPM2AncestralDiscrete' : KDPM2AncestralDiscreteScheduler,
'KDPM2Discrete' : KDPM2DiscreteScheduler,
'UniPCMultistep' : UniPCMultistepScheduler
}
def token_auto_concat_embeds(pipe, positive, negative):
max_length = pipe.tokenizer.model_max_length
positive_length = pipe.tokenizer(positive, return_tensors="pt").input_ids.shape[-1]
negative_length = pipe.tokenizer(negative, return_tensors="pt").input_ids.shape[-1]
print(f'Token length is model maximum: {max_length}, positive length: {positive_length}, negative length: {negative_length}.')
if max_length < positive_length or max_length < negative_length:
print('Concatenated embedding.')
if positive_length > negative_length:
positive_ids = pipe.tokenizer(positive, return_tensors="pt").input_ids.to("cuda")
negative_ids = pipe.tokenizer(negative, truncation=False, padding="max_length", max_length=positive_ids.shape[-1], return_tensors="pt").input_ids.to("cuda")
else:
negative_ids = pipe.tokenizer(negative, return_tensors="pt").input_ids.to("cuda")
positive_ids = pipe.tokenizer(positive, truncation=False, padding="max_length", max_length=negative_ids.shape[-1], return_tensors="pt").input_ids.to("cuda")
else:
positive_ids = pipe.tokenizer(positive, truncation=False, padding="max_length", max_length=max_length, return_tensors="pt").input_ids.to("cuda")
negative_ids = pipe.tokenizer(negative, truncation=False, padding="max_length", max_length=max_length, return_tensors="pt").input_ids.to("cuda")
positive_concat_embeds = []
negative_concat_embeds = []
for i in range(0, positive_ids.shape[-1], max_length):
positive_concat_embeds.append(pipe.text_encoder(positive_ids[:, i: i + max_length])[0])
negative_concat_embeds.append(pipe.text_encoder(negative_ids[:, i: i + max_length])[0])
positive_prompt_embeds = torch.cat(positive_concat_embeds, dim=1)
negative_prompt_embeds = torch.cat(negative_concat_embeds, dim=1)
return positive_prompt_embeds, negative_prompt_embeds
# Reference from : https://github.com/huggingface/diffusers/blob/main/scripts/convert_vae_pt_to_diffusers.py
def custom_convert_ldm_vae_checkpoint(checkpoint, config):
vae_state_dict = checkpoint
new_checkpoint = {}
new_checkpoint["encoder.conv_in.weight"] = vae_state_dict["encoder.conv_in.weight"]
new_checkpoint["encoder.conv_in.bias"] = vae_state_dict["encoder.conv_in.bias"]
new_checkpoint["encoder.conv_out.weight"] = vae_state_dict["encoder.conv_out.weight"]
new_checkpoint["encoder.conv_out.bias"] = vae_state_dict["encoder.conv_out.bias"]
new_checkpoint["encoder.conv_norm_out.weight"] = vae_state_dict["encoder.norm_out.weight"]
new_checkpoint["encoder.conv_norm_out.bias"] = vae_state_dict["encoder.norm_out.bias"]
new_checkpoint["decoder.conv_in.weight"] = vae_state_dict["decoder.conv_in.weight"]
new_checkpoint["decoder.conv_in.bias"] = vae_state_dict["decoder.conv_in.bias"]
new_checkpoint["decoder.conv_out.weight"] = vae_state_dict["decoder.conv_out.weight"]
new_checkpoint["decoder.conv_out.bias"] = vae_state_dict["decoder.conv_out.bias"]
new_checkpoint["decoder.conv_norm_out.weight"] = vae_state_dict["decoder.norm_out.weight"]
new_checkpoint["decoder.conv_norm_out.bias"] = vae_state_dict["decoder.norm_out.bias"]
new_checkpoint["quant_conv.weight"] = vae_state_dict["quant_conv.weight"]
new_checkpoint["quant_conv.bias"] = vae_state_dict["quant_conv.bias"]
new_checkpoint["post_quant_conv.weight"] = vae_state_dict["post_quant_conv.weight"]
new_checkpoint["post_quant_conv.bias"] = vae_state_dict["post_quant_conv.bias"]
# Retrieves the keys for the encoder down blocks only
num_down_blocks = len({".".join(layer.split(".")[:3]) for layer in vae_state_dict if "encoder.down" in layer})
down_blocks = {
layer_id: [key for key in vae_state_dict if f"down.{layer_id}" in key] for layer_id in range(num_down_blocks)
}
# Retrieves the keys for the decoder up blocks only
num_up_blocks = len({".".join(layer.split(".")[:3]) for layer in vae_state_dict if "decoder.up" in layer})
up_blocks = {
layer_id: [key for key in vae_state_dict if f"up.{layer_id}" in key] for layer_id in range(num_up_blocks)
}
for i in range(num_down_blocks):
resnets = [key for key in down_blocks[i] if f"down.{i}" in key and f"down.{i}.downsample" not in key]
if f"encoder.down.{i}.downsample.conv.weight" in vae_state_dict:
new_checkpoint[f"encoder.down_blocks.{i}.downsamplers.0.conv.weight"] = vae_state_dict.pop(
f"encoder.down.{i}.downsample.conv.weight"
)
new_checkpoint[f"encoder.down_blocks.{i}.downsamplers.0.conv.bias"] = vae_state_dict.pop(
f"encoder.down.{i}.downsample.conv.bias"
)
paths = renew_vae_resnet_paths(resnets)
meta_path = {"old": f"down.{i}.block", "new": f"down_blocks.{i}.resnets"}
assign_to_checkpoint(paths, new_checkpoint, vae_state_dict, additional_replacements=[meta_path], config=config)
mid_resnets = [key for key in vae_state_dict if "encoder.mid.block" in key]
num_mid_res_blocks = 2
for i in range(1, num_mid_res_blocks + 1):
resnets = [key for key in mid_resnets if f"encoder.mid.block_{i}" in key]
paths = renew_vae_resnet_paths(resnets)
meta_path = {"old": f"mid.block_{i}", "new": f"mid_block.resnets.{i - 1}"}
assign_to_checkpoint(paths, new_checkpoint, vae_state_dict, additional_replacements=[meta_path], config=config)
mid_attentions = [key for key in vae_state_dict if "encoder.mid.attn" in key]
paths = renew_vae_attention_paths(mid_attentions)
meta_path = {"old": "mid.attn_1", "new": "mid_block.attentions.0"}
assign_to_checkpoint(paths, new_checkpoint, vae_state_dict, additional_replacements=[meta_path], config=config)
conv_attn_to_linear(new_checkpoint)
for i in range(num_up_blocks):
block_id = num_up_blocks - 1 - i
resnets = [
key for key in up_blocks[block_id] if f"up.{block_id}" in key and f"up.{block_id}.upsample" not in key
]
if f"decoder.up.{block_id}.upsample.conv.weight" in vae_state_dict:
new_checkpoint[f"decoder.up_blocks.{i}.upsamplers.0.conv.weight"] = vae_state_dict[
f"decoder.up.{block_id}.upsample.conv.weight"
]
new_checkpoint[f"decoder.up_blocks.{i}.upsamplers.0.conv.bias"] = vae_state_dict[
f"decoder.up.{block_id}.upsample.conv.bias"
]
paths = renew_vae_resnet_paths(resnets)
meta_path = {"old": f"up.{block_id}.block", "new": f"up_blocks.{i}.resnets"}
assign_to_checkpoint(paths, new_checkpoint, vae_state_dict, additional_replacements=[meta_path], config=config)
mid_resnets = [key for key in vae_state_dict if "decoder.mid.block" in key]
num_mid_res_blocks = 2
for i in range(1, num_mid_res_blocks + 1):
resnets = [key for key in mid_resnets if f"decoder.mid.block_{i}" in key]
paths = renew_vae_resnet_paths(resnets)
meta_path = {"old": f"mid.block_{i}", "new": f"mid_block.resnets.{i - 1}"}
assign_to_checkpoint(paths, new_checkpoint, vae_state_dict, additional_replacements=[meta_path], config=config)
mid_attentions = [key for key in vae_state_dict if "decoder.mid.attn" in key]
paths = renew_vae_attention_paths(mid_attentions)
meta_path = {"old": "mid.attn_1", "new": "mid_block.attentions.0"}
assign_to_checkpoint(paths, new_checkpoint, vae_state_dict, additional_replacements=[meta_path], config=config)
conv_attn_to_linear(new_checkpoint)
return new_checkpoint
# Reference from : https://github.com/huggingface/diffusers/blob/main/scripts/convert_vae_pt_to_diffusers.py
def vae_pt_to_vae_diffuser(
checkpoint_path: str,
output_path: str,
):
# Only support V1
r = requests.get(
" https://raw.githubusercontent.com/CompVis/stable-diffusion/main/configs/stable-diffusion/v1-inference.yaml"
)
io_obj = io.BytesIO(r.content)
original_config = OmegaConf.load(io_obj)
image_size = 512
device = "cuda" if torch.cuda.is_available() else "cpu"
if checkpoint_path.endswith("safetensors"):
from safetensors import safe_open
checkpoint = {}
with safe_open(checkpoint_path, framework="pt", device="cpu") as f:
for key in f.keys():
checkpoint[key] = f.get_tensor(key)
else:
checkpoint = torch.load(checkpoint_path, map_location=device)["state_dict"]
# Convert the VAE model.
vae_config = create_vae_diffusers_config(original_config, image_size=image_size)
converted_vae_checkpoint = custom_convert_ldm_vae_checkpoint(checkpoint, vae_config)
vae = AutoencoderKL(**vae_config)
vae.load_state_dict(converted_vae_checkpoint)
vae.save_pretrained(output_path)
def convert_images_to_tensors(images: list[Image.Image]):
return torch.stack([np.transpose(ToTensor()(image), (1, 2, 0)) for image in images])
def convert_tensors_to_images(images: torch.tensor):
return [Image.fromarray(np.clip(255. * image.cpu().numpy(), 0, 255).astype(np.uint8)) for image in images]
def resize_images(images: list[Image.Image], size: tuple[int, int]):
return [image.resize(size) for image in images]