├── .github └── workflows │ └── publish.yml ├── .gitignore ├── LICENSE ├── README.md ├── __init__.py ├── img ├── workflow.png ├── workflow2.png ├── workflow3.png └── workflow3demo.gif ├── nodes.py ├── pyproject.toml ├── requirements.txt └── utils.py /.github/workflows/publish.yml: -------------------------------------------------------------------------------- 1 | name: Publish to Comfy registry 2 | on: 3 | workflow_dispatch: 4 | push: 5 | branches: 6 | - main 7 | paths: 8 | - "pyproject.toml" 9 | 10 | jobs: 11 | publish-node: 12 | name: Publish Custom Node to registry 13 | runs-on: ubuntu-latest 14 | steps: 15 | - name: Check out code 16 | uses: actions/checkout@v4 17 | - name: Publish Custom Node 18 | uses: Comfy-Org/publish-node-action@main 19 | with: 20 | ## Add your own personal access token to your Github Repository secrets and reference it here. 21 | personal_access_token: ${{ secrets.REGISTRY_ACCESS_TOKEN }} 22 | -------------------------------------------------------------------------------- /.gitignore: -------------------------------------------------------------------------------- 1 | # Byte-compiled / optimized / DLL files 2 | __pycache__/ 3 | *.py[cod] 4 | *$py.class 5 | 6 | # C extensions 7 | *.so 8 | 9 | # Distribution / packaging 10 | .Python 11 | build/ 12 | develop-eggs/ 13 | dist/ 14 | downloads/ 15 | eggs/ 16 | .eggs/ 17 | lib/ 18 | lib64/ 19 | parts/ 20 | sdist/ 21 | var/ 22 | wheels/ 23 | share/python-wheels/ 24 | *.egg-info/ 25 | .installed.cfg 26 | *.egg 27 | MANIFEST 28 | 29 | # PyInstaller 30 | # Usually these files are written by a python script from a template 31 | # before PyInstaller builds the exe, so as to inject date/other infos into it. 32 | *.manifest 33 | *.spec 34 | 35 | # Installer logs 36 | pip-log.txt 37 | pip-delete-this-directory.txt 38 | 39 | # Unit test / coverage reports 40 | htmlcov/ 41 | .tox/ 42 | .nox/ 43 | .coverage 44 | .coverage.* 45 | .cache 46 | nosetests.xml 47 | coverage.xml 48 | *.cover 49 | *.py,cover 50 | .hypothesis/ 51 | .pytest_cache/ 52 | cover/ 53 | 54 | # Translations 55 | *.mo 56 | *.pot 57 | 58 | # Django stuff: 59 | *.log 60 | local_settings.py 61 | db.sqlite3 62 | db.sqlite3-journal 63 | 64 | # Flask stuff: 65 | instance/ 66 | .webassets-cache 67 | 68 | # Scrapy stuff: 69 | .scrapy 70 | 71 | # Sphinx documentation 72 | docs/_build/ 73 | 74 | # PyBuilder 75 | .pybuilder/ 76 | target/ 77 | 78 | # Jupyter Notebook 79 | .ipynb_checkpoints 80 | 81 | # IPython 82 | profile_default/ 83 | ipython_config.py 84 | 85 | # pyenv 86 | # For a library or package, you might want to ignore these files since the code is 87 | # intended to run in multiple environments; otherwise, check them in: 88 | # .python-version 89 | 90 | # pipenv 91 | # According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control. 92 | # However, in case of collaboration, if having platform-specific dependencies or dependencies 93 | # having no cross-platform support, pipenv may install dependencies that don't work, or not 94 | # install all needed dependencies. 95 | #Pipfile.lock 96 | 97 | # poetry 98 | # Similar to Pipfile.lock, it is generally recommended to include poetry.lock in version control. 99 | # This is especially recommended for binary packages to ensure reproducibility, and is more 100 | # commonly ignored for libraries. 101 | # https://python-poetry.org/docs/basic-usage/#commit-your-poetrylock-file-to-version-control 102 | #poetry.lock 103 | 104 | # pdm 105 | # Similar to Pipfile.lock, it is generally recommended to include pdm.lock in version control. 106 | #pdm.lock 107 | # pdm stores project-wide configurations in .pdm.toml, but it is recommended to not include it 108 | # in version control. 109 | # https://pdm.fming.dev/#use-with-ide 110 | .pdm.toml 111 | 112 | # PEP 582; used by e.g. github.com/David-OConnor/pyflow and github.com/pdm-project/pdm 113 | __pypackages__/ 114 | 115 | # Celery stuff 116 | celerybeat-schedule 117 | celerybeat.pid 118 | 119 | # SageMath parsed files 120 | *.sage.py 121 | 122 | # Environments 123 | .env 124 | .venv 125 | env/ 126 | venv/ 127 | ENV/ 128 | env.bak/ 129 | venv.bak/ 130 | 131 | # Spyder project settings 132 | .spyderproject 133 | .spyproject 134 | 135 | # Rope project settings 136 | .ropeproject 137 | 138 | # mkdocs documentation 139 | /site 140 | 141 | # mypy 142 | .mypy_cache/ 143 | .dmypy.json 144 | dmypy.json 145 | 146 | # Pyre type checker 147 | .pyre/ 148 | 149 | # pytype static type analyzer 150 | .pytype/ 151 | 152 | # Cython debug symbols 153 | cython_debug/ 154 | 155 | # PyCharm 156 | # JetBrains specific template is maintained in a separate JetBrains.gitignore that can 157 | # be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore 158 | # and can be added to the global gitignore or merged into this file. For a more nuclear 159 | # option (not recommended) you can uncomment the following to ignore the entire idea folder. 160 | #.idea/ 161 | 162 | StreamDiffusion/* -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | MIT License 2 | 3 | Copyright (c) 2023 Limitex 4 | 5 | Permission is hereby granted, free of charge, to any person obtaining a copy 6 | of this software and associated documentation files (the "Software"), to deal 7 | in the Software without restriction, including without limitation the rights 8 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 9 | copies of the Software, and to permit persons to whom the Software is 10 | furnished to do so, subject to the following conditions: 11 | 12 | The above copyright notice and this permission notice shall be included in all 13 | copies or substantial portions of the Software. 14 | 15 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 16 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 17 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 18 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 19 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 20 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE 21 | SOFTWARE. 22 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # ComfyUI-Diffusers 2 | 3 | This repository is a custom node in ComfyUI. 4 | 5 | ## Overview 6 | 7 | ### Workflow 1 8 | 9 | This is a program that allows you to use Hugging Face Diffusers module with ComfyUI. Additionally, Stream Diffusion is also available. 10 | 11 | ![Workflow](img/workflow.png) 12 | 13 | ### Workflow 2 14 | 15 | In addition, real-time generation is possible by doing the following. 16 | When running, please enable Auto Queue in Extra options. 17 | 18 | ![Workflow2](img/workflow2.png) 19 | 20 | ### Workflow 3 21 | 22 | In combination with [VideoHelperSuite](https://github.com/Kosinkadink/ComfyUI-VideoHelperSuite.git), you can also run vid2vid. 23 | 24 | ![Workflow3](img/workflow3.png) 25 | 26 | The execution looks like this: 27 | 28 | ![Workflow3demo](img/workflow3demo.gif) 29 | 30 | 31 | ## Usage 32 | 33 | Run the following command inside ComfyUI/custom_nodes. 34 | 35 | ```cmd 36 | git clone https://github.com/Limitex/ComfyUI-Diffusers.git 37 | cd ComfyUI-Diffusers 38 | pip install -r requirements.txt 39 | git clone https://github.com/cumulo-autumn/StreamDiffusion.git 40 | python -m streamdiffusion.tools.install-tensorrt 41 | ``` 42 | ## Recommended Custom Nodes 43 | 44 | https://github.com/Kosinkadink/ComfyUI-VideoHelperSuite.git 45 | 46 | ## Node 47 | 48 | ### Diffusers Pipeline Loader (DiffusersPipelineLoader) 49 | 50 | ### Diffusers Vae Loader (DiffusersVaeLoader) 51 | 52 | ### Diffusers Scheduler Loader (DiffusersSchedulerLoader) 53 | 54 | ### Diffusers Model Makeup (DiffusersModelMakeup) 55 | 56 | ### Diffusers Clip Text Encode (DiffusersClipTextEncode) 57 | 58 | ### Diffusers Sampler (DiffusersSampler) 59 | 60 | ### Create Int List (CreateIntListNode) 61 | 62 | ### LcmLoraLoader (LcmLoraLoader) 63 | 64 | ### StreamDiffusion Create Stream (StreamDiffusionCreateStream) 65 | 66 | ### StreamDiffusion Sampler (StreamDiffusionSampler) 67 | 68 | ### StreamDiffusion Warmup (StreamDiffusionWarmup) 69 | 70 | ### StreamDiffusion Fast Sampler (StreamDiffusionFastSampler) 71 | 72 | ## Reference 73 | 74 | https://github.com/cumulo-autumn/StreamDiffusion 75 | -------------------------------------------------------------------------------- /__init__.py: -------------------------------------------------------------------------------- 1 | from .nodes import * 2 | 3 | __all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS"] 4 | -------------------------------------------------------------------------------- /img/workflow.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Limitex/ComfyUI-Diffusers/4498af1394797d2725ffb6a4142ff3e4cd0d35a3/img/workflow.png -------------------------------------------------------------------------------- /img/workflow2.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Limitex/ComfyUI-Diffusers/4498af1394797d2725ffb6a4142ff3e4cd0d35a3/img/workflow2.png -------------------------------------------------------------------------------- /img/workflow3.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Limitex/ComfyUI-Diffusers/4498af1394797d2725ffb6a4142ff3e4cd0d35a3/img/workflow3.png -------------------------------------------------------------------------------- /img/workflow3demo.gif: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Limitex/ComfyUI-Diffusers/4498af1394797d2725ffb6a4142ff3e4cd0d35a3/img/workflow3demo.gif -------------------------------------------------------------------------------- /nodes.py: -------------------------------------------------------------------------------- 1 | import copy 2 | import os 3 | import torch 4 | from safetensors.torch import load_file 5 | from .utils import SCHEDULERS, token_auto_concat_embeds, vae_pt_to_vae_diffuser, convert_images_to_tensors, convert_tensors_to_images, resize_images 6 | from comfy.model_management import get_torch_device 7 | import folder_paths 8 | from streamdiffusion import StreamDiffusion 9 | from streamdiffusion.image_utils import postprocess_image 10 | from diffusers import StableDiffusionPipeline, AutoencoderKL, AutoencoderTiny 11 | 12 | 13 | class DiffusersPipelineLoader: 14 | def __init__(self): 15 | self.tmp_dir = folder_paths.get_temp_directory() 16 | self.dtype = torch.float32 17 | 18 | @classmethod 19 | def INPUT_TYPES(s): 20 | return {"required": { "ckpt_name": (folder_paths.get_filename_list("checkpoints"), ), }} 21 | 22 | RETURN_TYPES = ("PIPELINE", "AUTOENCODER", "SCHEDULER",) 23 | 24 | FUNCTION = "create_pipeline" 25 | 26 | CATEGORY = "Diffusers" 27 | 28 | def create_pipeline(self, ckpt_name): 29 | ckpt_cache_path = os.path.join(self.tmp_dir, ckpt_name) 30 | 31 | StableDiffusionPipeline.from_single_file( 32 | pretrained_model_link_or_path=folder_paths.get_full_path("checkpoints", ckpt_name), 33 | torch_dtype=self.dtype, 34 | cache_dir=self.tmp_dir, 35 | ).save_pretrained(ckpt_cache_path, safe_serialization=True) 36 | 37 | pipe = StableDiffusionPipeline.from_pretrained( 38 | pretrained_model_name_or_path=ckpt_cache_path, 39 | torch_dtype=self.dtype, 40 | cache_dir=self.tmp_dir, 41 | ) 42 | return ((pipe, ckpt_cache_path), pipe.vae, pipe.scheduler) 43 | 44 | class DiffusersVaeLoader: 45 | def __init__(self): 46 | self.tmp_dir = folder_paths.get_temp_directory() 47 | self.dtype = torch.float32 48 | 49 | @classmethod 50 | def INPUT_TYPES(s): 51 | return {"required": { "vae_name": (folder_paths.get_filename_list("vae"), ), }} 52 | 53 | RETURN_TYPES = ("AUTOENCODER",) 54 | 55 | FUNCTION = "create_pipeline" 56 | 57 | CATEGORY = "Diffusers" 58 | 59 | def create_pipeline(self, vae_name): 60 | ckpt_cache_path = os.path.join(self.tmp_dir, vae_name) 61 | vae_pt_to_vae_diffuser(folder_paths.get_full_path("vae", vae_name), ckpt_cache_path) 62 | 63 | vae = AutoencoderKL.from_pretrained( 64 | pretrained_model_name_or_path=ckpt_cache_path, 65 | torch_dtype=self.dtype, 66 | cache_dir=self.tmp_dir, 67 | ) 68 | 69 | return (vae,) 70 | 71 | class DiffusersSchedulerLoader: 72 | def __init__(self): 73 | self.tmp_dir = folder_paths.get_temp_directory() 74 | self.dtype = torch.float32 75 | 76 | @classmethod 77 | def INPUT_TYPES(s): 78 | return { 79 | "required": { 80 | "pipeline": ("PIPELINE", ), 81 | "scheduler_name": (list(SCHEDULERS.keys()), ), 82 | } 83 | } 84 | 85 | RETURN_TYPES = ("SCHEDULER",) 86 | 87 | FUNCTION = "load_scheduler" 88 | 89 | CATEGORY = "Diffusers" 90 | 91 | def load_scheduler(self, pipeline, scheduler_name): 92 | scheduler = SCHEDULERS[scheduler_name].from_pretrained( 93 | pretrained_model_name_or_path=pipeline[1], 94 | torch_dtype=self.dtype, 95 | cache_dir=self.tmp_dir, 96 | subfolder='scheduler' 97 | ) 98 | return (scheduler,) 99 | 100 | class DiffusersModelMakeup: 101 | def __init__(self): 102 | self.torch_device = get_torch_device() 103 | 104 | @classmethod 105 | def INPUT_TYPES(s): 106 | return { 107 | "required": { 108 | "pipeline": ("PIPELINE", ), 109 | "scheduler": ("SCHEDULER", ), 110 | "autoencoder": ("AUTOENCODER", ), 111 | }, 112 | } 113 | 114 | RETURN_TYPES = ("MAKED_PIPELINE",) 115 | 116 | FUNCTION = "makeup_pipeline" 117 | 118 | CATEGORY = "Diffusers" 119 | 120 | def makeup_pipeline(self, pipeline, scheduler, autoencoder): 121 | pipeline = pipeline[0] 122 | pipeline.vae = autoencoder 123 | pipeline.scheduler = scheduler 124 | pipeline.safety_checker = None if pipeline.safety_checker is None else lambda images, **kwargs: (images, [False]) 125 | pipeline.enable_attention_slicing() 126 | pipeline = pipeline.to(self.torch_device) 127 | return (pipeline,) 128 | 129 | class DiffusersClipTextEncode: 130 | @classmethod 131 | def INPUT_TYPES(s): 132 | return {"required": { 133 | "maked_pipeline": ("MAKED_PIPELINE", ), 134 | "positive": ("STRING", {"multiline": True}), 135 | "negative": ("STRING", {"multiline": True}), 136 | }} 137 | 138 | RETURN_TYPES = ("EMBEDS", "EMBEDS", "STRING", "STRING", ) 139 | RETURN_NAMES = ("positive_embeds", "negative_embeds", "positive", "negative", ) 140 | 141 | FUNCTION = "concat_embeds" 142 | 143 | CATEGORY = "Diffusers" 144 | 145 | def concat_embeds(self, maked_pipeline, positive, negative): 146 | positive_embeds, negative_embeds = token_auto_concat_embeds(maked_pipeline, positive,negative) 147 | 148 | return (positive_embeds, negative_embeds, positive, negative, ) 149 | 150 | class DiffusersSampler: 151 | def __init__(self): 152 | self.torch_device = get_torch_device() 153 | 154 | @classmethod 155 | def INPUT_TYPES(s): 156 | return {"required": { 157 | "maked_pipeline": ("MAKED_PIPELINE", ), 158 | "positive_embeds": ("EMBEDS", ), 159 | "negative_embeds": ("EMBEDS", ), 160 | "width": ("INT", {"default": 512, "min": 1, "max": 8192, "step": 1}), 161 | "height": ("INT", {"default": 512, "min": 1, "max": 8192, "step": 1}), 162 | "steps": ("INT", {"default": 20, "min": 1, "max": 10000}), 163 | "cfg": ("FLOAT", {"default": 8.0, "min": 0.0, "max": 100.0, "step":0.1, "round": 0.01}), 164 | "seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff}), 165 | }} 166 | 167 | RETURN_TYPES = ("IMAGE",) 168 | 169 | FUNCTION = "sample" 170 | 171 | CATEGORY = "Diffusers" 172 | 173 | def sample(self, maked_pipeline, positive_embeds, negative_embeds, height, width, steps, cfg, seed): 174 | images = maked_pipeline( 175 | prompt_embeds=positive_embeds, 176 | height=height, 177 | width=width, 178 | num_inference_steps=steps, 179 | guidance_scale=cfg, 180 | negative_prompt_embeds=negative_embeds, 181 | generator=torch.Generator(self.torch_device).manual_seed(seed) 182 | ).images 183 | return (convert_images_to_tensors(images),) 184 | 185 | # - Stream Diffusion - 186 | 187 | class CreateIntListNode: 188 | @classmethod 189 | def INPUT_TYPES(s): 190 | max_element = 10 191 | return { 192 | "required": { 193 | "elements_count" : ("INT", {"default": 2, "min": 1, "max": max_element, "step": 1}), 194 | }, 195 | "optional": { 196 | f"element_{i}": ("INT", {"default": 0}) for i in range(1, max_element) 197 | } 198 | } 199 | 200 | RETURN_TYPES = ("LIST",) 201 | FUNCTION = "create_list" 202 | 203 | CATEGORY = "Diffusers/StreamDiffusion" 204 | 205 | def create_list(self, elements_count, **kwargs): 206 | return ([value for key, value in kwargs.items()][:elements_count], ) 207 | 208 | class LcmLoraLoader: 209 | @classmethod 210 | def INPUT_TYPES(s): 211 | return {"required": { "lora_name": (folder_paths.get_filename_list("loras"), ), }} 212 | 213 | RETURN_TYPES = ("LCM_LORA",) 214 | FUNCTION = "load_lora" 215 | 216 | CATEGORY = "Diffusers/StreamDiffusion" 217 | 218 | def load_lora(self, lora_name): 219 | return (load_file(folder_paths.get_full_path("loras", lora_name)), ) 220 | 221 | class StreamDiffusionCreateStream: 222 | def __init__(self): 223 | self.dtype = torch.float32 224 | self.torch_device = get_torch_device() 225 | self.tmp_dir = folder_paths.get_temp_directory() 226 | 227 | @classmethod 228 | def INPUT_TYPES(s): 229 | return { 230 | "required": { 231 | "maked_pipeline": ("MAKED_PIPELINE", ), 232 | "t_index_list": ("LIST", ), 233 | "width": ("INT", {"default": 512, "min": 1, "max": 8192, "step": 1}), 234 | "height": ("INT", {"default": 512, "min": 1, "max": 8192, "step": 1}), 235 | "do_add_noise": ("BOOLEAN", {"default": True}), 236 | "use_denoising_batch": ("BOOLEAN", {"default": True}), 237 | "frame_buffer_size": ("INT", {"default": 1, "min": 1, "max": 10000}), 238 | "cfg_type": (["none", "full", "self", "initialize"], {"default": "none"}), 239 | "xformers_memory_efficient_attention": ("BOOLEAN", {"default": False}), 240 | "lcm_lora" : ("LCM_LORA", ), 241 | "tiny_vae" : ("STRING", {"default": "madebyollin/taesd"}) 242 | }, 243 | } 244 | 245 | RETURN_TYPES = ("STREAM",) 246 | FUNCTION = "load_stream" 247 | 248 | CATEGORY = "Diffusers/StreamDiffusion" 249 | 250 | 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): 251 | maked_pipeline = copy.deepcopy(maked_pipeline) 252 | lcm_lora = copy.deepcopy(lcm_lora) 253 | stream = StreamDiffusion( 254 | pipe = maked_pipeline, 255 | t_index_list = t_index_list, 256 | torch_dtype = self.dtype, 257 | width = width, 258 | height = height, 259 | do_add_noise = do_add_noise, 260 | use_denoising_batch = use_denoising_batch, 261 | frame_buffer_size = frame_buffer_size, 262 | cfg_type = cfg_type, 263 | ) 264 | stream.load_lcm_lora(lcm_lora) 265 | stream.fuse_lora() 266 | stream.vae = AutoencoderTiny.from_pretrained( 267 | pretrained_model_name_or_path=tiny_vae, 268 | torch_dtype=self.dtype, 269 | cache_dir=self.tmp_dir, 270 | ).to( 271 | device=maked_pipeline.device, 272 | dtype=maked_pipeline.dtype 273 | ) 274 | 275 | if xformers_memory_efficient_attention: 276 | maked_pipeline.enable_xformers_memory_efficient_attention() 277 | return (stream, ) 278 | 279 | class StreamDiffusionSampler: 280 | def __init__(self): 281 | self.torch_device = get_torch_device() 282 | 283 | @classmethod 284 | def INPUT_TYPES(s): 285 | return { 286 | "required": { 287 | "stream": ("STREAM", ), 288 | "positive": ("STRING", {"multiline": True}), 289 | "negative": ("STRING", {"multiline": True}), 290 | "steps": ("INT", {"default": 50, "min": 1, "max": 10000}), 291 | "cfg": ("FLOAT", {"default": 1.2, "min": 0.0, "max": 100.0}), 292 | "delta": ("FLOAT", {"default": 1, "min": 0.0, "max": 1.0}), 293 | "seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff}), 294 | "num": ("INT", {"default": 1, "min": 1, "max": 10000}), 295 | "warmup": ("INT", {"default": 1, "min": 0, "max": 10000}), 296 | }, 297 | "optional" : { 298 | "image" : ("IMAGE", ) 299 | } 300 | } 301 | 302 | RETURN_TYPES = ("IMAGE",) 303 | 304 | FUNCTION = "sample" 305 | 306 | CATEGORY = "Diffusers/StreamDiffusion" 307 | 308 | def sample(self, stream: StreamDiffusion, positive, negative, steps, cfg, delta, seed, num, warmup, image = None): 309 | stream.prepare( 310 | prompt = positive, 311 | negative_prompt = negative, 312 | num_inference_steps = steps, 313 | guidance_scale = cfg, 314 | delta = delta, 315 | seed = seed 316 | ) 317 | 318 | if image != None: 319 | image = convert_tensors_to_images(image) 320 | image = resize_images(image, (stream.width, stream.height)) 321 | 322 | for _ in range(warmup): 323 | stream() 324 | 325 | result = [] 326 | for _ in range(num): 327 | x_outputs = [] 328 | if image is None: 329 | x_outputs.append(stream.txt2img()) 330 | else: 331 | stream(image[0]) 332 | for i in image[1:] + image[-1:]: 333 | x_outputs.append(stream(i)) 334 | for x_output in x_outputs: 335 | result.append(postprocess_image(x_output, output_type="pil")[0]) 336 | 337 | return (convert_images_to_tensors(result),) 338 | 339 | class StreamDiffusionWarmup: 340 | @classmethod 341 | def INPUT_TYPES(s): 342 | return { 343 | "required": { 344 | "stream": ("STREAM", ), 345 | "negative": ("STRING", {"multiline": True}), 346 | "steps": ("INT", {"default": 50, "min": 1, "max": 10000}), 347 | "cfg": ("FLOAT", {"default": 1.2, "min": 0.0, "max": 100.0}), 348 | "delta": ("FLOAT", {"default": 1, "min": 0.0, "max": 1.0}), 349 | "seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff}), 350 | "warmup": ("INT", {"default": 1, "min": 0, "max": 10000}), 351 | }, 352 | } 353 | 354 | RETURN_TYPES = ("WARMUP_STREAM",) 355 | 356 | FUNCTION = "stream_warmup" 357 | 358 | CATEGORY = "Diffusers/StreamDiffusion" 359 | 360 | def stream_warmup(self, stream: StreamDiffusion, negative, steps, cfg, delta, seed, warmup): 361 | stream.prepare( 362 | prompt="", 363 | negative_prompt=negative, 364 | num_inference_steps = steps, 365 | guidance_scale = cfg, 366 | delta = delta, 367 | seed = seed 368 | ) 369 | 370 | for _ in range(warmup): 371 | stream() 372 | 373 | return (stream, ) 374 | 375 | 376 | class StreamDiffusionFastSampler: 377 | @classmethod 378 | def INPUT_TYPES(s): 379 | return { 380 | "required": { 381 | "warmup_stream": ("WARMUP_STREAM", ), 382 | "positive": ("STRING", {"multiline": True}), 383 | "num": ("INT", {"default": 1, "min": 1, "max": 10000}), 384 | }, 385 | } 386 | 387 | RETURN_TYPES = ("IMAGE",) 388 | 389 | FUNCTION = "sample" 390 | 391 | CATEGORY = "Diffusers/StreamDiffusion" 392 | 393 | def sample(self, warmup_stream, positive, num): 394 | stream: StreamDiffusion = warmup_stream 395 | 396 | stream.update_prompt(positive) 397 | 398 | result = [] 399 | for _ in range(num): 400 | x_output = stream.txt2img() 401 | result.append(postprocess_image(x_output, output_type="pil")[0]) 402 | return (convert_images_to_tensors(result),) 403 | 404 | # - - - - - - - - - - - - - - - - - - 405 | 406 | 407 | NODE_CLASS_MAPPINGS = { 408 | "DiffusersPipelineLoader": DiffusersPipelineLoader, 409 | "DiffusersVaeLoader": DiffusersVaeLoader, 410 | "DiffusersSchedulerLoader": DiffusersSchedulerLoader, 411 | "DiffusersModelMakeup": DiffusersModelMakeup, 412 | "DiffusersClipTextEncode": DiffusersClipTextEncode, 413 | "DiffusersSampler": DiffusersSampler, 414 | "CreateIntListNode": CreateIntListNode, 415 | "LcmLoraLoader": LcmLoraLoader, 416 | "StreamDiffusionCreateStream": StreamDiffusionCreateStream, 417 | "StreamDiffusionSampler": StreamDiffusionSampler, 418 | "StreamDiffusionWarmup": StreamDiffusionWarmup, 419 | "StreamDiffusionFastSampler": StreamDiffusionFastSampler, 420 | } 421 | 422 | NODE_DISPLAY_NAME_MAPPINGS = { 423 | "DiffusersPipelineLoader": "Diffusers Pipeline Loader", 424 | "DiffusersVaeLoader": "Diffusers Vae Loader", 425 | "DiffusersSchedulerLoader": "Diffusers Scheduler Loader", 426 | "DiffusersModelMakeup": "Diffusers Model Makeup", 427 | "DiffusersClipTextEncode": "Diffusers Clip Text Encode", 428 | "DiffusersSampler": "Diffusers Sampler", 429 | "CreateIntListNode": "Create Int List", 430 | "LcmLoraLoader": "LCM Lora Loader", 431 | "StreamDiffusionCreateStream": "StreamDiffusion Create Stream", 432 | "StreamDiffusionSampler": "StreamDiffusion Sampler", 433 | "StreamDiffusionWarmup": "StreamDiffusion Warmup", 434 | "StreamDiffusionFastSampler": "StreamDiffusion Fast Sampler", 435 | } 436 | -------------------------------------------------------------------------------- /pyproject.toml: -------------------------------------------------------------------------------- 1 | [project] 2 | name = "comfyui-diffusers" 3 | description = "This extension enables the use of the diffuser pipeline in ComfyUI. It also includes nodes related to Stream Diffusion." 4 | version = "1.0.2" 5 | license = { file = "LICENSE" } 6 | dependencies = ["diffusers[torch]>=0.29.0", "accelerate", "transformers", "safetensors", "omegaconf", "pytorch_lightning", "xformers", "git+https://github.com/cumulo-autumn/StreamDiffusion.git@main#egg=streamdiffusion[tensorrt]"] 7 | 8 | [project.urls] 9 | Repository = "https://github.com/Limitex/ComfyUI-Diffusers" 10 | # Used by Comfy Registry https://comfyregistry.org 11 | 12 | [tool.comfy] 13 | PublisherId = "limitex" 14 | DisplayName = "ComfyUI-Diffusers" 15 | Icon = "" 16 | -------------------------------------------------------------------------------- /requirements.txt: -------------------------------------------------------------------------------- 1 | diffusers[torch]>=0.29.0 2 | accelerate 3 | transformers 4 | safetensors 5 | omegaconf 6 | pytorch_lightning 7 | xformers 8 | git+https://github.com/cumulo-autumn/StreamDiffusion.git@main#egg=streamdiffusion[tensorrt] 9 | -------------------------------------------------------------------------------- /utils.py: -------------------------------------------------------------------------------- 1 | import io 2 | import torch 3 | import requests 4 | import numpy as np 5 | from PIL import Image 6 | from omegaconf import OmegaConf 7 | from torchvision.transforms import ToTensor 8 | from diffusers.pipelines.stable_diffusion.convert_from_ckpt import ( 9 | assign_to_checkpoint, 10 | conv_attn_to_linear, 11 | create_vae_diffusers_config, 12 | renew_vae_attention_paths, 13 | renew_vae_resnet_paths, 14 | ) 15 | from diffusers import ( 16 | AutoencoderKL, 17 | DDIMScheduler, 18 | DDPMScheduler, 19 | DEISMultistepScheduler, 20 | DPMSolverMultistepScheduler, 21 | DPMSolverSinglestepScheduler, 22 | EulerAncestralDiscreteScheduler, 23 | EulerDiscreteScheduler, 24 | HeunDiscreteScheduler, 25 | KDPM2AncestralDiscreteScheduler, 26 | KDPM2DiscreteScheduler, 27 | UniPCMultistepScheduler, 28 | ) 29 | 30 | SCHEDULERS = { 31 | 'DDIM' : DDIMScheduler, 32 | 'DDPM' : DDPMScheduler, 33 | 'DEISMultistep' : DEISMultistepScheduler, 34 | 'DPMSolverMultistep' : DPMSolverMultistepScheduler, 35 | 'DPMSolverSinglestep' : DPMSolverSinglestepScheduler, 36 | 'EulerAncestralDiscrete' : EulerAncestralDiscreteScheduler, 37 | 'EulerDiscrete' : EulerDiscreteScheduler, 38 | 'HeunDiscrete' : HeunDiscreteScheduler, 39 | 'KDPM2AncestralDiscrete' : KDPM2AncestralDiscreteScheduler, 40 | 'KDPM2Discrete' : KDPM2DiscreteScheduler, 41 | 'UniPCMultistep' : UniPCMultistepScheduler 42 | } 43 | 44 | def token_auto_concat_embeds(pipe, positive, negative): 45 | max_length = pipe.tokenizer.model_max_length 46 | positive_length = pipe.tokenizer(positive, return_tensors="pt").input_ids.shape[-1] 47 | negative_length = pipe.tokenizer(negative, return_tensors="pt").input_ids.shape[-1] 48 | 49 | print(f'Token length is model maximum: {max_length}, positive length: {positive_length}, negative length: {negative_length}.') 50 | if max_length < positive_length or max_length < negative_length: 51 | print('Concatenated embedding.') 52 | if positive_length > negative_length: 53 | positive_ids = pipe.tokenizer(positive, return_tensors="pt").input_ids.to("cuda") 54 | negative_ids = pipe.tokenizer(negative, truncation=False, padding="max_length", max_length=positive_ids.shape[-1], return_tensors="pt").input_ids.to("cuda") 55 | else: 56 | negative_ids = pipe.tokenizer(negative, return_tensors="pt").input_ids.to("cuda") 57 | positive_ids = pipe.tokenizer(positive, truncation=False, padding="max_length", max_length=negative_ids.shape[-1], return_tensors="pt").input_ids.to("cuda") 58 | else: 59 | positive_ids = pipe.tokenizer(positive, truncation=False, padding="max_length", max_length=max_length, return_tensors="pt").input_ids.to("cuda") 60 | negative_ids = pipe.tokenizer(negative, truncation=False, padding="max_length", max_length=max_length, return_tensors="pt").input_ids.to("cuda") 61 | 62 | positive_concat_embeds = [] 63 | negative_concat_embeds = [] 64 | for i in range(0, positive_ids.shape[-1], max_length): 65 | positive_concat_embeds.append(pipe.text_encoder(positive_ids[:, i: i + max_length])[0]) 66 | negative_concat_embeds.append(pipe.text_encoder(negative_ids[:, i: i + max_length])[0]) 67 | 68 | positive_prompt_embeds = torch.cat(positive_concat_embeds, dim=1) 69 | negative_prompt_embeds = torch.cat(negative_concat_embeds, dim=1) 70 | return positive_prompt_embeds, negative_prompt_embeds 71 | 72 | # Reference from : https://github.com/huggingface/diffusers/blob/main/scripts/convert_vae_pt_to_diffusers.py 73 | def custom_convert_ldm_vae_checkpoint(checkpoint, config): 74 | vae_state_dict = checkpoint 75 | 76 | new_checkpoint = {} 77 | 78 | new_checkpoint["encoder.conv_in.weight"] = vae_state_dict["encoder.conv_in.weight"] 79 | new_checkpoint["encoder.conv_in.bias"] = vae_state_dict["encoder.conv_in.bias"] 80 | new_checkpoint["encoder.conv_out.weight"] = vae_state_dict["encoder.conv_out.weight"] 81 | new_checkpoint["encoder.conv_out.bias"] = vae_state_dict["encoder.conv_out.bias"] 82 | new_checkpoint["encoder.conv_norm_out.weight"] = vae_state_dict["encoder.norm_out.weight"] 83 | new_checkpoint["encoder.conv_norm_out.bias"] = vae_state_dict["encoder.norm_out.bias"] 84 | 85 | new_checkpoint["decoder.conv_in.weight"] = vae_state_dict["decoder.conv_in.weight"] 86 | new_checkpoint["decoder.conv_in.bias"] = vae_state_dict["decoder.conv_in.bias"] 87 | new_checkpoint["decoder.conv_out.weight"] = vae_state_dict["decoder.conv_out.weight"] 88 | new_checkpoint["decoder.conv_out.bias"] = vae_state_dict["decoder.conv_out.bias"] 89 | new_checkpoint["decoder.conv_norm_out.weight"] = vae_state_dict["decoder.norm_out.weight"] 90 | new_checkpoint["decoder.conv_norm_out.bias"] = vae_state_dict["decoder.norm_out.bias"] 91 | 92 | new_checkpoint["quant_conv.weight"] = vae_state_dict["quant_conv.weight"] 93 | new_checkpoint["quant_conv.bias"] = vae_state_dict["quant_conv.bias"] 94 | new_checkpoint["post_quant_conv.weight"] = vae_state_dict["post_quant_conv.weight"] 95 | new_checkpoint["post_quant_conv.bias"] = vae_state_dict["post_quant_conv.bias"] 96 | 97 | # Retrieves the keys for the encoder down blocks only 98 | num_down_blocks = len({".".join(layer.split(".")[:3]) for layer in vae_state_dict if "encoder.down" in layer}) 99 | down_blocks = { 100 | layer_id: [key for key in vae_state_dict if f"down.{layer_id}" in key] for layer_id in range(num_down_blocks) 101 | } 102 | 103 | # Retrieves the keys for the decoder up blocks only 104 | num_up_blocks = len({".".join(layer.split(".")[:3]) for layer in vae_state_dict if "decoder.up" in layer}) 105 | up_blocks = { 106 | layer_id: [key for key in vae_state_dict if f"up.{layer_id}" in key] for layer_id in range(num_up_blocks) 107 | } 108 | 109 | for i in range(num_down_blocks): 110 | resnets = [key for key in down_blocks[i] if f"down.{i}" in key and f"down.{i}.downsample" not in key] 111 | 112 | if f"encoder.down.{i}.downsample.conv.weight" in vae_state_dict: 113 | new_checkpoint[f"encoder.down_blocks.{i}.downsamplers.0.conv.weight"] = vae_state_dict.pop( 114 | f"encoder.down.{i}.downsample.conv.weight" 115 | ) 116 | new_checkpoint[f"encoder.down_blocks.{i}.downsamplers.0.conv.bias"] = vae_state_dict.pop( 117 | f"encoder.down.{i}.downsample.conv.bias" 118 | ) 119 | 120 | paths = renew_vae_resnet_paths(resnets) 121 | meta_path = {"old": f"down.{i}.block", "new": f"down_blocks.{i}.resnets"} 122 | assign_to_checkpoint(paths, new_checkpoint, vae_state_dict, additional_replacements=[meta_path], config=config) 123 | 124 | mid_resnets = [key for key in vae_state_dict if "encoder.mid.block" in key] 125 | num_mid_res_blocks = 2 126 | for i in range(1, num_mid_res_blocks + 1): 127 | resnets = [key for key in mid_resnets if f"encoder.mid.block_{i}" in key] 128 | 129 | paths = renew_vae_resnet_paths(resnets) 130 | meta_path = {"old": f"mid.block_{i}", "new": f"mid_block.resnets.{i - 1}"} 131 | assign_to_checkpoint(paths, new_checkpoint, vae_state_dict, additional_replacements=[meta_path], config=config) 132 | 133 | mid_attentions = [key for key in vae_state_dict if "encoder.mid.attn" in key] 134 | paths = renew_vae_attention_paths(mid_attentions) 135 | meta_path = {"old": "mid.attn_1", "new": "mid_block.attentions.0"} 136 | assign_to_checkpoint(paths, new_checkpoint, vae_state_dict, additional_replacements=[meta_path], config=config) 137 | conv_attn_to_linear(new_checkpoint) 138 | 139 | for i in range(num_up_blocks): 140 | block_id = num_up_blocks - 1 - i 141 | resnets = [ 142 | key for key in up_blocks[block_id] if f"up.{block_id}" in key and f"up.{block_id}.upsample" not in key 143 | ] 144 | 145 | if f"decoder.up.{block_id}.upsample.conv.weight" in vae_state_dict: 146 | new_checkpoint[f"decoder.up_blocks.{i}.upsamplers.0.conv.weight"] = vae_state_dict[ 147 | f"decoder.up.{block_id}.upsample.conv.weight" 148 | ] 149 | new_checkpoint[f"decoder.up_blocks.{i}.upsamplers.0.conv.bias"] = vae_state_dict[ 150 | f"decoder.up.{block_id}.upsample.conv.bias" 151 | ] 152 | 153 | paths = renew_vae_resnet_paths(resnets) 154 | meta_path = {"old": f"up.{block_id}.block", "new": f"up_blocks.{i}.resnets"} 155 | assign_to_checkpoint(paths, new_checkpoint, vae_state_dict, additional_replacements=[meta_path], config=config) 156 | 157 | mid_resnets = [key for key in vae_state_dict if "decoder.mid.block" in key] 158 | num_mid_res_blocks = 2 159 | for i in range(1, num_mid_res_blocks + 1): 160 | resnets = [key for key in mid_resnets if f"decoder.mid.block_{i}" in key] 161 | 162 | paths = renew_vae_resnet_paths(resnets) 163 | meta_path = {"old": f"mid.block_{i}", "new": f"mid_block.resnets.{i - 1}"} 164 | assign_to_checkpoint(paths, new_checkpoint, vae_state_dict, additional_replacements=[meta_path], config=config) 165 | 166 | mid_attentions = [key for key in vae_state_dict if "decoder.mid.attn" in key] 167 | paths = renew_vae_attention_paths(mid_attentions) 168 | meta_path = {"old": "mid.attn_1", "new": "mid_block.attentions.0"} 169 | assign_to_checkpoint(paths, new_checkpoint, vae_state_dict, additional_replacements=[meta_path], config=config) 170 | conv_attn_to_linear(new_checkpoint) 171 | return new_checkpoint 172 | 173 | # Reference from : https://github.com/huggingface/diffusers/blob/main/scripts/convert_vae_pt_to_diffusers.py 174 | def vae_pt_to_vae_diffuser( 175 | checkpoint_path: str, 176 | output_path: str, 177 | ): 178 | # Only support V1 179 | r = requests.get( 180 | " https://raw.githubusercontent.com/CompVis/stable-diffusion/main/configs/stable-diffusion/v1-inference.yaml" 181 | ) 182 | io_obj = io.BytesIO(r.content) 183 | 184 | original_config = OmegaConf.load(io_obj) 185 | image_size = 512 186 | device = "cuda" if torch.cuda.is_available() else "cpu" 187 | if checkpoint_path.endswith("safetensors"): 188 | from safetensors import safe_open 189 | 190 | checkpoint = {} 191 | with safe_open(checkpoint_path, framework="pt", device="cpu") as f: 192 | for key in f.keys(): 193 | checkpoint[key] = f.get_tensor(key) 194 | else: 195 | checkpoint = torch.load(checkpoint_path, map_location=device)["state_dict"] 196 | 197 | # Convert the VAE model. 198 | vae_config = create_vae_diffusers_config(original_config, image_size=image_size) 199 | converted_vae_checkpoint = custom_convert_ldm_vae_checkpoint(checkpoint, vae_config) 200 | 201 | vae = AutoencoderKL(**vae_config) 202 | vae.load_state_dict(converted_vae_checkpoint) 203 | vae.save_pretrained(output_path) 204 | 205 | 206 | def convert_images_to_tensors(images: list[Image.Image]): 207 | return torch.stack([np.transpose(ToTensor()(image), (1, 2, 0)) for image in images]) 208 | 209 | def convert_tensors_to_images(images: torch.tensor): 210 | return [Image.fromarray(np.clip(255. * image.cpu().numpy(), 0, 255).astype(np.uint8)) for image in images] 211 | 212 | def resize_images(images: list[Image.Image], size: tuple[int, int]): 213 | return [image.resize(size) for image in images] --------------------------------------------------------------------------------