├── .gitattributes ├── .gitignore ├── README.md ├── __init__.py ├── nodes.py └── requirements.txt /.gitattributes: -------------------------------------------------------------------------------- 1 | # Auto detect text files and perform LF normalization 2 | * text=auto 3 | -------------------------------------------------------------------------------- /.gitignore: -------------------------------------------------------------------------------- 1 | __pycache__ 2 | /venv 3 | /checkpoints 4 | .vscode 5 | *.ckpt 6 | *.safetensors 7 | *.pth 8 | types -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # ComfyUI StableCascade using diffusers 2 | 3 | # Update: This repo is already deprecated as the native support is in ComfyUI, more info: https://gist.github.com/comfyanonymous/0f09119a342d0dd825bb2d99d19b781c 4 | 5 | Simple quick wrapper for https://huggingface.co/stabilityai/stable-cascade 6 | 7 | Comfy is going to implement this properly soon, this repo is just for quick testing for the impatient! 8 | 9 | Currently requires this diffusers branch: 10 | `pip install git+https://github.com/kashif/diffusers.git@wuerstchen-v3` 11 | -------------------------------------------------------------------------------- /__init__.py: -------------------------------------------------------------------------------- 1 | from .nodes import NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS 2 | 3 | __all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS"] -------------------------------------------------------------------------------- /nodes.py: -------------------------------------------------------------------------------- 1 | import os 2 | import torch 3 | from torchvision.transforms import ToTensor 4 | 5 | from diffusers import StableCascadeDecoderPipeline, StableCascadePriorPipeline 6 | import comfy.model_management 7 | 8 | script_directory = os.path.dirname(os.path.abspath(__file__)) 9 | 10 | class DiffusersStableCascade: 11 | @classmethod 12 | def INPUT_TYPES(s): 13 | return {"required": { 14 | "width": ("INT", {"default": 512, "min": 128, "max": 8192, "step": 128}), 15 | "height": ("INT", {"default": 512, "min": 128, "max": 8192, "step": 128}), 16 | "seed": ("INT", {"default": 123,"min": 0, "max": 0xffffffffffffffff, "step": 1}), 17 | "guidance_scale": ("FLOAT", {"default": 4.0, "min": 0.01, "max": 100.0, "step": 0.01}), 18 | "steps": ("INT", {"default": 20, "min": 1, "max": 4096, "step": 1}), 19 | "decoder_steps": ("INT", {"default": 10, "min": 1, "max": 4096, "step": 1}), 20 | "batch_size": ("INT", {"default": 1, "min": 1, "max": 4096, "step": 1}), 21 | "prompt": ("STRING", {"multiline": True, "default": "",}), 22 | "negative_prompt": ("STRING", {"multiline": True, "default": "",}), 23 | "model_cpu_offload": ("BOOLEAN", {"default": False}), 24 | }, 25 | 26 | } 27 | 28 | RETURN_TYPES = ("IMAGE",) 29 | RETURN_NAMES =("image",) 30 | FUNCTION = "process" 31 | 32 | CATEGORY = "DiffusersStableCascade" 33 | 34 | def process(self, width, height, seed, steps, guidance_scale, prompt, negative_prompt, batch_size, decoder_steps, model_cpu_offload): 35 | 36 | comfy.model_management.unload_all_models() 37 | torch.manual_seed(seed) 38 | 39 | device = comfy.model_management.get_torch_device() 40 | 41 | if not hasattr(self, 'prior') or not hasattr(self, 'decoder'): 42 | 43 | self.prior = StableCascadePriorPipeline.from_pretrained("stabilityai/stable-cascade-prior", torch_dtype=torch.bfloat16).to(device) 44 | self.decoder = StableCascadeDecoderPipeline.from_pretrained("stabilityai/stable-cascade", torch_dtype=torch.float16).to(device) 45 | if model_cpu_offload: 46 | self.prior.enable_model_cpu_offload() 47 | self.decoder.enable_model_cpu_offload() 48 | 49 | prior_output = self.prior( 50 | prompt=prompt, 51 | height=height, 52 | width=width, 53 | negative_prompt=negative_prompt, 54 | guidance_scale=guidance_scale, 55 | num_images_per_prompt=batch_size, 56 | num_inference_steps=steps 57 | ) 58 | decoder_output = self.decoder( 59 | image_embeddings=prior_output.image_embeddings.half(), 60 | prompt=prompt, 61 | negative_prompt=negative_prompt, 62 | guidance_scale=0.0, 63 | output_type="pil", 64 | num_inference_steps=decoder_steps 65 | ).images 66 | 67 | tensors = [ToTensor()(img) for img in decoder_output] 68 | batch_tensor = torch.stack(tensors).permute(0, 2, 3, 1).cpu() 69 | 70 | 71 | return (batch_tensor,) 72 | 73 | 74 | NODE_CLASS_MAPPINGS = { 75 | "DiffusersStableCascade": DiffusersStableCascade, 76 | } 77 | NODE_DISPLAY_NAME_MAPPINGS = { 78 | "DiffusersStableCascade": "DiffusersStableCascade", 79 | } -------------------------------------------------------------------------------- /requirements.txt: -------------------------------------------------------------------------------- 1 | -e git+https://github.com/kashif/diffusers.git@wuerstchen-v3#egg=diffusers 2 | typing-extensions>=4.9.0 3 | accelerate --------------------------------------------------------------------------------