diff --git a/RAUNET.md b/RAUNET.md new file mode 100644 index 0000000..a2bf3fa --- /dev/null +++ b/RAUNET.md @@ -0,0 +1,29 @@ +During investigation of compatibility issues with [WASasquatch's FreeU_Advanced](https://github.com/WASasquatch/FreeU_Advanced/tree/main) and [blepping's jank HiDiffusion](https://github.com/blepping/comfyui_jankhidiffusion) nodes I stumbled upon some quite hard problems. There are `FreeU` nodes in ComfyUI, but no such for HiDiffusion, so I decided to implement RAUNet on base of my BrushNet implementation. blepping, I am sorry. :) + +What is RAUNet? I know many of you saw and generate images with a lot of limbs, fingers and faces all morphed together. + +The authors of HiDiffusion invent simple, yet efficient trick to alleviate this problem. Here is an example: + +![example workflow](example/RAUNet1.png?raw=true) + +[workflow](example/RAUNet_basic.json) + +The left picture is created using ZavyChromaXL checkpoint on 2048x2048 canvas. The right one using RAUNet. + +In my experience the node is helpful but quite sensitive to its parameters. And there is no universal solution - you should adjust them for every new image you generate. It also lowers model's creativity, you usually get only what you described in prompt Look at the example: in first you have a forest in the background, but RAUNet deleted all except fox which is described in the prompt. + +There are two independent parts in this node: DU (Downsample/Upsample) and XA (CrossAttension). The four parameters are the start and end steps for applying these parts. + +The Downsample/Upsample part lowers models degrees of freedom. If you apply it a lot (for more steps) the resulting images will have a lot of symmetries. + +The CrossAttension part lowers number of objects which model tracks in image. + +Usually you apply DU and after several steps apply XA, sometimes you will need only XA, you should try it yourself. + +This is ControlNet example. The lower image is pure model, the upper is after using RAUNet. You can see small fox and two tails in lower image. + +![example workflow](example/RAUNet2.png?raw=true) + +[workflow](example/RAUNet_with_CN.json) + +The node can be implemented for any model. Right now it can be applyed to SD15 and SDXL models. \ No newline at end of file diff --git a/README.md b/README.md index aca71a0..0d6284b 100644 --- a/README.md +++ b/README.md @@ -1,11 +1,17 @@ ## ComfyUI-BrushNet -Custom nodes for ComfyUI allow to inpaint using Brushnet: ["BrushNet: A Plug-and-Play Image Inpainting Model with Decomposed Dual-Branch Diffusion"](https://arxiv.org/abs/2403.06976). +These are custom nodes for ComfyUI native implementation of -My contribution is limited to the ComfyUI adaptation, and all credit goes to the authors of the paper. +- Brushnet: ["BrushNet: A Plug-and-Play Image Inpainting Model with Decomposed Dual-Branch Diffusion"](https://arxiv.org/abs/2403.06976) +- PowerPaint: [A Task is Worth One Word: Learning with Task Prompts for High-Quality Versatile Image Inpainting](https://arxiv.org/abs/2312.03594) +- HiDiffusion: [HiDiffusion: Unlocking Higher-Resolution Creativity and Efficiency in Pretrained Diffusion Models](https://arxiv.org/abs/2311.17528) + +My contribution is limited to the ComfyUI adaptation, and all credit goes to the authors of the papers. ## Updates +May 16, 2024. Internal rework to improve compatibility with other nodes. [RAUNet](RAUNET.md) is implemented. + May 12, 2024. CutForInpaint node, see [example](BIG_IMAGE.md). May 11, 2024. Image batch is implemented. You can even add BrushNet to AnimateDiff vid2vid workflow, but they don't work together - they are different models and both try to patch UNet. Added some more examples. @@ -43,7 +49,7 @@ The checkpoint in `segmentation_mask_brushnet_ckpt` provides checkpoints trained `segmentation_mask_brushnet_ckpt` and `random_mask_brushnet_ckpt` contains BrushNet for SD 1.5 models while `segmentation_mask_brushnet_ckpt_sdxl_v0` and `random_mask_brushnet_ckpt_sdxl_v0` for SDXL. -You should place `diffusion_pytorch_model.safetensors` files to your `models/inpaint` folder. +You should place `diffusion_pytorch_model.safetensors` files to your `models/inpaint` folder. You can also specify `inpaint` folder in your `extra_model_paths.yaml`. For PowerPaint you should download three files. Both `diffusion_pytorch_model.safetensors` and `pytorch_model.bin` from [here](https://huggingface.co/JunhaoZhuang/PowerPaint_v2/tree/main/PowerPaint_Brushnet) should be placed in your `models/inpaint` folder. @@ -220,12 +226,11 @@ List of known uncompartible nodes. - [WASasquatch's FreeU_Advanced](https://github.com/WASasquatch/FreeU_Advanced/tree/main) - [blepping's jank HiDiffusion](https://github.com/blepping/comfyui_jankhidiffusion) -I will think how to avoid it. - ## Credits The code is based on - [BrushNet](https://github.com/TencentARC/BrushNet) - [PowerPaint](https://github.com/zhuang2002/PowerPaint) +- [HiDiffusion](https://github.com/megvii-research/HiDiffusion) - [diffusers](https://github.com/huggingface/diffusers) diff --git a/__init__.py b/__init__.py index 9f6459d..ef5d143 100644 --- a/__init__.py +++ b/__init__.py @@ -1,4 +1,5 @@ from .brushnet_nodes import BrushNetLoader, BrushNet, BlendInpaint, PowerPaintCLIPLoader, PowerPaint, CutForInpaint +from .raunet_nodes import RAUNet """ @author: nullquant @title: BrushNet @@ -15,6 +16,7 @@ NODE_CLASS_MAPPINGS = { "PowerPaintCLIPLoader": PowerPaintCLIPLoader, "PowerPaint": PowerPaint, "CutForInpaint": CutForInpaint, + "RAUNet": RAUNet, } # A dictionary that contains the friendly/humanly readable titles for the nodes @@ -25,4 +27,5 @@ NODE_DISPLAY_NAME_MAPPINGS = { "PowerPaintCLIPLoader": "PowerPaint CLIP Loader", "PowerPaint": "PowerPaint", "CutForInpaint": "Cut For Inpaint", + "RAUNet": "RAUNet", } diff --git a/brushnet_nodes.py b/brushnet_nodes.py index 7a15687..64f89a4 100644 --- a/brushnet_nodes.py +++ b/brushnet_nodes.py @@ -1,19 +1,20 @@ +import os +import types +from typing import Tuple + import torch import torchvision.transforms as T import torch.nn.functional as F -#import math - -import os -import folder_paths +from accelerate import init_empty_weights, load_checkpoint_and_dispatch #import sys #from sys import platform # Get the parent directory of 'comfy' and add it to the Python path #comfy_parent_dir = os.path.abspath(os.path.join(os.path.dirname(__file__), '../../')) #sys.path.append(comfy_parent_dir) -import comfy -from accelerate import init_empty_weights, load_checkpoint_and_dispatch +import comfy +import folder_paths from .model_patch import add_model_patch_option, patch_model_function_wrapper @@ -22,12 +23,6 @@ from .brushnet.brushnet_ca import BrushNetModel as PowerPaintModel from .brushnet.powerpaint_utils import TokenizerWrapper, add_tokens -import types -from typing import Tuple -#import warnings -#warnings.filterwarnings("ignore", category=UserWarning, module="torch") -#warnings.filterwarnings("ignore", category=UserWarning, module="safetensors") - current_directory = os.path.dirname(os.path.abspath(__file__)) brushnet_config_file = os.path.join(current_directory, 'brushnet', 'brushnet.json') brushnet_xl_config_file = os.path.join(current_directory, 'brushnet', 'brushnet_xl.json') @@ -36,12 +31,6 @@ powerpaint_config_file = os.path.join(current_directory,'brushnet', 'powerpaint. sd15_scaling_factor = 0.18215 sdxl_scaling_factor = 0.13025 -#ComfySDLayers = [comfy.ops.disable_weight_init.Conv2d, -# comfy.ldm.modules.attention.SpatialTransformer, -# comfy.ldm.modules.diffusionmodules.openaimodel.Downsample, -# comfy.ldm.modules.diffusionmodules.openaimodel.ResBlock -# ] - ModelsToUnload = [comfy.sd1_clip.SD1ClipModel, comfy.ldm.models.autoencoder.AutoencoderKL ] @@ -587,7 +576,9 @@ def get_files_with_extension(folder_name, extension=['safetensors']): for ext in extension: if ext in name: abs_list.append(os.path.join(x[0], name)) - abs_list = set(abs_list) + + abs_list = sorted(list(set(abs_list))) + names = [] for x in abs_list: remain = x @@ -995,7 +986,7 @@ def add_brushnet_patch(model, brushnet, torch_dtype, conditioning_latents, return -1, layer_list.reverse() return len(layer_list) - 1 - layer_list.index(tp), layer_list - def brushnet_forward(model, x, timesteps, transformer_options): + def brushnet_forward(model, x, timesteps, transformer_options, control): if 'brushnet' not in transformer_options['model_patch']: input_samples = [] mid_sample = 0 @@ -1060,8 +1051,12 @@ def add_brushnet_patch(model, brushnet, torch_dtype, conditioning_latents, def forward_patched_by_brushnet(self, x, *args, **kwargs): h = self.original_forward(x, *args, **kwargs) if hasattr(self, 'add_sample_after') and type(self): - if torch.is_tensor(self.add_sample_after): - h += self.add_sample_after.to(h.dtype).to(h.device) + to_add = self.add_sample_after + if torch.is_tensor(to_add): + # interpolate due to RAUNet + if h.shape[2] != to_add.shape[2] or h.shape[3] != to_add.shape[3]: + to_add = torch.nn.functional.interpolate(to_add, size=(h.shape[2], h.shape[3]), mode='bicubic') + h += to_add.to(h.dtype).to(h.device) else: h += self.add_sample_after self.add_sample_after = 0 diff --git a/example/RAUNet1.png b/example/RAUNet1.png new file mode 100644 index 0000000..9b0f86c Binary files /dev/null and b/example/RAUNet1.png differ diff --git a/example/RAUNet2.png b/example/RAUNet2.png new file mode 100644 index 0000000..d01c844 Binary files /dev/null and b/example/RAUNet2.png differ diff --git a/example/RAUNet_basic.json b/example/RAUNet_basic.json new file mode 100644 index 0000000..ef2c07d --- /dev/null +++ b/example/RAUNet_basic.json @@ -0,0 +1 @@ +{"last_node_id": 26, "last_link_id": 48, "nodes": [{"id": 7, "type": "KSamplerAdvanced", "pos": [1281, 461], "size": {"0": 315, "1": 334}, "flags": {}, "order": 5, "mode": 0, "inputs": [{"name": "model", "type": "MODEL", "link": 48}, {"name": "positive", "type": "CONDITIONING", "link": 8}, {"name": "negative", "type": "CONDITIONING", "link": 9}, {"name": "latent_image", "type": "LATENT", "link": 10, "slot_index": 3}], "outputs": [{"name": "LATENT", "type": "LATENT", "links": [22], "shape": 3, "slot_index": 0}], "properties": {"Node name for S&R": "KSamplerAdvanced"}, "widgets_values": ["disable", 0, "fixed", 25, 8, "ddpm", "normal", 0, 10000, "disable"], "color": "#2a363b", "bgcolor": "#3f5159"}, {"id": 23, "type": "KSamplerAdvanced", "pos": [1280, 872], "size": {"0": 315, "1": 334}, "flags": {}, "order": 6, "mode": 0, "inputs": [{"name": "model", "type": "MODEL", "link": 46}, {"name": "positive", "type": "CONDITIONING", "link": 45}, {"name": "negative", "type": "CONDITIONING", "link": 43}, {"name": "latent_image", "type": "LATENT", "link": 44, "slot_index": 3}], "outputs": [{"name": "LATENT", "type": "LATENT", "links": [42], "shape": 3, "slot_index": 0}], "properties": {"Node name for S&R": "KSamplerAdvanced"}, "widgets_values": ["disable", 0, "fixed", 25, 8, "ddpm", "normal", 0, 10000, "disable"], "color": "#2a363b", "bgcolor": "#3f5159"}, {"id": 1, "type": "CheckpointLoaderSimple", "pos": [452, 461], "size": {"0": 320.2000732421875, "1": 108.99996948242188}, "flags": {}, "order": 0, "mode": 0, "outputs": [{"name": "MODEL", "type": "MODEL", "links": [46, 47], "shape": 3, "slot_index": 0}, {"name": "CLIP", "type": "CLIP", "links": [1, 6], "shape": 3, "slot_index": 1}, {"name": "VAE", "type": "VAE", "links": [23, 40], "shape": 3, "slot_index": 2}], "title": "Load Base Checkpoint", "properties": {"Node name for S&R": "CheckpointLoaderSimple"}, "widgets_values": ["SDXL/zavychromaxl_v70.safetensors"], "color": "#2a363b", "bgcolor": "#3f5159"}, {"id": 5, "type": "CLIPTextEncodeSDXL", "pos": [854, 844], "size": [319.27423095703125, 311.4324369430542], "flags": {}, "order": 4, "mode": 0, "inputs": [{"name": "clip", "type": "CLIP", "link": 6}], "outputs": [{"name": "CONDITIONING", "type": "CONDITIONING", "links": [9, 43], "shape": 3, "slot_index": 0}], "properties": {"Node name for S&R": "CLIPTextEncodeSDXL"}, "widgets_values": [4096, 4096, 0, 0, 1024, 1024, "ugly, deformed, noisy, low poly, blurry, text, duplicate, poorly drawn, mosaic", "ugly, deformed, noisy, low poly, blurry, text, duplicate, poorly drawn, mosaic"], "color": "#322", "bgcolor": "#533"}, {"id": 8, "type": "EmptyLatentImage", "pos": [868, 1205], "size": [269.2342041015627, 106], "flags": {}, "order": 1, "mode": 0, "outputs": [{"name": "LATENT", "type": "LATENT", "links": [10, 44], "shape": 3, "slot_index": 0}], "properties": {"Node name for S&R": "EmptyLatentImage"}, "widgets_values": [2048, 2048, 1]}, {"id": 2, "type": "CLIPTextEncodeSDXL", "pos": [852, 457], "size": [325.67423095703134, 332.3523832321167], "flags": {}, "order": 3, "mode": 0, "inputs": [{"name": "clip", "type": "CLIP", "link": 1}], "outputs": [{"name": "CONDITIONING", "type": "CONDITIONING", "links": [8, 45], "shape": 3, "slot_index": 0}], "properties": {"Node name for S&R": "CLIPTextEncodeSDXL"}, "widgets_values": [4096, 4096, 0, 0, 1024, 1024, "an exotic fox, cute, chibi kawaii. detailed fur, hyperdetailed, big reflective eyes", "thick strokes, bright colors, fairytale, artstation,centered composition, perfect composition, centered, vibrant colors, muted colors, high detailed, 8k"], "color": "#232", "bgcolor": "#353"}, {"id": 25, "type": "PreviewImage", "pos": [1628, 460], "size": {"0": 650.7540893554688, "1": 766.8323974609375}, "flags": {}, "order": 10, "mode": 0, "inputs": [{"name": "images", "type": "IMAGE", "link": 41}], "properties": {"Node name for S&R": "PreviewImage"}}, {"id": 24, "type": "VAEDecode", "pos": [1683, 350], "size": {"0": 210, "1": 46}, "flags": {}, "order": 8, "mode": 0, "inputs": [{"name": "samples", "type": "LATENT", "link": 42, "slot_index": 0}, {"name": "vae", "type": "VAE", "link": 40, "slot_index": 1}], "outputs": [{"name": "IMAGE", "type": "IMAGE", "links": [41], "shape": 3, "slot_index": 0}], "properties": {"Node name for S&R": "VAEDecode"}}, {"id": 14, "type": "VAEDecode", "pos": [1985, 347], "size": {"0": 210, "1": 46}, "flags": {}, "order": 7, "mode": 0, "inputs": [{"name": "samples", "type": "LATENT", "link": 22}, {"name": "vae", "type": "VAE", "link": 23}], "outputs": [{"name": "IMAGE", "type": "IMAGE", "links": [24], "shape": 3, "slot_index": 0}], "properties": {"Node name for S&R": "VAEDecode"}}, {"id": 15, "type": "PreviewImage", "pos": [2298, 457], "size": [650.7540725708009, 766.8323699951172], "flags": {}, "order": 9, "mode": 0, "inputs": [{"name": "images", "type": "IMAGE", "link": 24}], "properties": {"Node name for S&R": "PreviewImage"}}, {"id": 26, "type": "RAUNet", "pos": [857, 270], "size": {"0": 315, "1": 130}, "flags": {}, "order": 2, "mode": 0, "inputs": [{"name": "model", "type": "MODEL", "link": 47}], "outputs": [{"name": "model", "type": "MODEL", "links": [48], "shape": 3, "slot_index": 0}], "properties": {"Node name for S&R": "RAUNet"}, "widgets_values": [0, 2, 2, 6]}], "links": [[1, 1, 1, 2, 0, "CLIP"], [6, 1, 1, 5, 0, "CLIP"], [8, 2, 0, 7, 1, "CONDITIONING"], [9, 5, 0, 7, 2, "CONDITIONING"], [10, 8, 0, 7, 3, "LATENT"], [22, 7, 0, 14, 0, "LATENT"], [23, 1, 2, 14, 1, "VAE"], [24, 14, 0, 15, 0, "IMAGE"], [40, 1, 2, 24, 1, "VAE"], [41, 24, 0, 25, 0, "IMAGE"], [42, 23, 0, 24, 0, "LATENT"], [43, 5, 0, 23, 2, "CONDITIONING"], [44, 8, 0, 23, 3, "LATENT"], [45, 2, 0, 23, 1, "CONDITIONING"], [46, 1, 0, 23, 0, "MODEL"], [47, 1, 0, 26, 0, "MODEL"], [48, 26, 0, 7, 0, "MODEL"]], "groups": [], "config": {}, "extra": {}, "version": 0.4} \ No newline at end of file diff --git a/example/RAUNet_with_CN.json b/example/RAUNet_with_CN.json new file mode 100644 index 0000000..e7f97cc --- /dev/null +++ b/example/RAUNet_with_CN.json @@ -0,0 +1 @@ +{"last_node_id": 20, "last_link_id": 34, "nodes": [{"id": 5, "type": "VAEDecode", "pos": [1916.4395019531253, 183.40589904785156], "size": {"0": 210, "1": 46}, "flags": {}, "order": 12, "mode": 0, "inputs": [{"name": "samples", "type": "LATENT", "link": 6}, {"name": "vae", "type": "VAE", "link": 8, "slot_index": 1}], "outputs": [{"name": "IMAGE", "type": "IMAGE", "links": [7], "shape": 3, "slot_index": 0}], "properties": {"Node name for S&R": "VAEDecode"}}, {"id": 10, "type": "ControlNetLoader", "pos": [229, -378], "size": {"0": 432.609130859375, "1": 78.54664611816406}, "flags": {}, "order": 0, "mode": 0, "outputs": [{"name": "CONTROL_NET", "type": "CONTROL_NET", "links": [15], "shape": 3, "slot_index": 0}], "properties": {"Node name for S&R": "ControlNetLoader"}, "widgets_values": ["control_canny-fp16.safetensors"]}, {"id": 11, "type": "LoadImage", "pos": [230, -237], "size": {"0": 393.4891357421875, "1": 460.0666809082031}, "flags": {}, "order": 1, "mode": 0, "outputs": [{"name": "IMAGE", "type": "IMAGE", "links": [18], "shape": 3, "slot_index": 0}, {"name": "MASK", "type": "MASK", "links": null, "shape": 3}], "properties": {"Node name for S&R": "LoadImage"}, "widgets_values": ["fox_with_sword (1).png", "image"]}, {"id": 2, "type": "CLIPTextEncode", "pos": [713, 274], "size": {"0": 392.8395080566406, "1": 142.005859375}, "flags": {}, "order": 6, "mode": 0, "inputs": [{"name": "clip", "type": "CLIP", "link": 1}], "outputs": [{"name": "CONDITIONING", "type": "CONDITIONING", "links": [16], "shape": 3, "slot_index": 0}], "properties": {"Node name for S&R": "CLIPTextEncode"}, "widgets_values": ["black fox with big sword, standing in a small town cute, chibi kawaii. fairytale, artstation, centered composition, perfect composition, centered, vibrant colors, muted colors, high detailed, 8k"], "color": "#232", "bgcolor": "#353"}, {"id": 16, "type": "VAEDecode", "pos": [1897, -612], "size": {"0": 210, "1": 46}, "flags": {}, "order": 13, "mode": 0, "inputs": [{"name": "samples", "type": "LATENT", "link": 24}, {"name": "vae", "type": "VAE", "link": 26, "slot_index": 1}], "outputs": [{"name": "IMAGE", "type": "IMAGE", "links": [25], "shape": 3, "slot_index": 0}], "properties": {"Node name for S&R": "VAEDecode"}}, {"id": 9, "type": "ControlNetApply", "pos": [1153, 111], "size": {"0": 317.4000244140625, "1": 98}, "flags": {}, "order": 9, "mode": 0, "inputs": [{"name": "conditioning", "type": "CONDITIONING", "link": 16}, {"name": "control_net", "type": "CONTROL_NET", "link": 15}, {"name": "image", "type": "IMAGE", "link": 19}], "outputs": [{"name": "CONDITIONING", "type": "CONDITIONING", "links": [17, 27], "shape": 3, "slot_index": 0}], "properties": {"Node name for S&R": "ControlNetApply"}, "widgets_values": [0.8]}, {"id": 3, "type": "CLIPTextEncode", "pos": [714, 488], "size": {"0": 399.75958251953125, "1": 111.60586547851562}, "flags": {}, "order": 7, "mode": 0, "inputs": [{"name": "clip", "type": "CLIP", "link": 2}], "outputs": [{"name": "CONDITIONING", "type": "CONDITIONING", "links": [5, 28], "shape": 3, "slot_index": 0}], "properties": {"Node name for S&R": "CLIPTextEncode"}, "widgets_values": ["hat, text, blurry, ugly, duplicate, poorly drawn, deformed, mosaic"], "color": "#322", "bgcolor": "#533"}, {"id": 1, "type": "CheckpointLoaderSimple", "pos": [199, 295], "size": {"0": 440.8395080566406, "1": 99.80586242675781}, "flags": {}, "order": 2, "mode": 0, "outputs": [{"name": "MODEL", "type": "MODEL", "links": [23, 29], "shape": 3, "slot_index": 0}, {"name": "CLIP", "type": "CLIP", "links": [1, 2], "shape": 3, "slot_index": 1}, {"name": "VAE", "type": "VAE", "links": [8, 26], "shape": 3, "slot_index": 2}], "properties": {"Node name for S&R": "CheckpointLoaderSimple"}, "widgets_values": ["SD15/revAnimated_v2Rebirth.safetensors"]}, {"id": 13, "type": "PreviewImage", "pos": [1057, -320], "size": {"0": 210, "1": 246}, "flags": {}, "order": 8, "mode": 0, "inputs": [{"name": "images", "type": "IMAGE", "link": 20}], "properties": {"Node name for S&R": "PreviewImage"}}, {"id": 6, "type": "PreviewImage", "pos": [1883, 277], "size": {"0": 623.648193359375, "1": 645.5486450195312}, "flags": {}, "order": 14, "mode": 0, "inputs": [{"name": "images", "type": "IMAGE", "link": 7}], "properties": {"Node name for S&R": "PreviewImage"}}, {"id": 17, "type": "PreviewImage", "pos": [1897, -508], "size": {"0": 602.8720703125, "1": 630.0126953125}, "flags": {}, "order": 15, "mode": 0, "inputs": [{"name": "images", "type": "IMAGE", "link": 25}], "properties": {"Node name for S&R": "PreviewImage"}}, {"id": 12, "type": "CannyEdgePreprocessor", "pos": [689, -123], "size": {"0": 315, "1": 106}, "flags": {}, "order": 4, "mode": 0, "inputs": [{"name": "image", "type": "IMAGE", "link": 18}], "outputs": [{"name": "IMAGE", "type": "IMAGE", "links": [19, 20], "shape": 3, "slot_index": 0}], "properties": {"Node name for S&R": "CannyEdgePreprocessor"}, "widgets_values": [100, 200, 1024]}, {"id": 18, "type": "RAUNet", "pos": [1113, -500], "size": {"0": 315, "1": 130}, "flags": {}, "order": 5, "mode": 0, "inputs": [{"name": "model", "type": "MODEL", "link": 29}], "outputs": [{"name": "model", "type": "MODEL", "links": [30], "shape": 3, "slot_index": 0}], "properties": {"Node name for S&R": "RAUNet"}, "widgets_values": [0, 2, 2, 8]}, {"id": 4, "type": "KSampler", "pos": [1529, 115], "size": {"0": 315, "1": 262}, "flags": {}, "order": 10, "mode": 0, "inputs": [{"name": "model", "type": "MODEL", "link": 23}, {"name": "positive", "type": "CONDITIONING", "link": 17}, {"name": "negative", "type": "CONDITIONING", "link": 5}, {"name": "latent_image", "type": "LATENT", "link": 9, "slot_index": 3}], "outputs": [{"name": "LATENT", "type": "LATENT", "links": [6], "shape": 3, "slot_index": 0}], "properties": {"Node name for S&R": "KSampler"}, "widgets_values": [0, "fixed", 25, 8, "euler_ancestral", "normal", 1]}, {"id": 15, "type": "KSampler", "pos": [1523, -441], "size": {"0": 315, "1": 262}, "flags": {}, "order": 11, "mode": 0, "inputs": [{"name": "model", "type": "MODEL", "link": 30}, {"name": "positive", "type": "CONDITIONING", "link": 27}, {"name": "negative", "type": "CONDITIONING", "link": 28}, {"name": "latent_image", "type": "LATENT", "link": 31, "slot_index": 3}], "outputs": [{"name": "LATENT", "type": "LATENT", "links": [24], "shape": 3, "slot_index": 0}], "properties": {"Node name for S&R": "KSampler"}, "widgets_values": [0, "fixed", 25, 8, "euler_ancestral", "normal", 1]}, {"id": 7, "type": "EmptyLatentImage", "pos": [797, 652], "size": {"0": 315, "1": 106}, "flags": {}, "order": 3, "mode": 0, "outputs": [{"name": "LATENT", "type": "LATENT", "links": [9, 31], "shape": 3, "slot_index": 0}], "properties": {"Node name for S&R": "EmptyLatentImage"}, "widgets_values": [1024, 1024, 1]}], "links": [[1, 1, 1, 2, 0, "CLIP"], [2, 1, 1, 3, 0, "CLIP"], [5, 3, 0, 4, 2, "CONDITIONING"], [6, 4, 0, 5, 0, "LATENT"], [7, 5, 0, 6, 0, "IMAGE"], [8, 1, 2, 5, 1, "VAE"], [9, 7, 0, 4, 3, "LATENT"], [15, 10, 0, 9, 1, "CONTROL_NET"], [16, 2, 0, 9, 0, "CONDITIONING"], [17, 9, 0, 4, 1, "CONDITIONING"], [18, 11, 0, 12, 0, "IMAGE"], [19, 12, 0, 9, 2, "IMAGE"], [20, 12, 0, 13, 0, "IMAGE"], [23, 1, 0, 4, 0, "MODEL"], [24, 15, 0, 16, 0, "LATENT"], [25, 16, 0, 17, 0, "IMAGE"], [26, 1, 2, 16, 1, "VAE"], [27, 9, 0, 15, 1, "CONDITIONING"], [28, 3, 0, 15, 2, "CONDITIONING"], [29, 1, 0, 18, 0, "MODEL"], [30, 18, 0, 15, 0, "MODEL"], [31, 7, 0, 15, 3, "LATENT"]], "groups": [], "config": {}, "extra": {}, "version": 0.4} \ No newline at end of file diff --git a/model_patch.py b/model_patch.py index 76c7387..2c7e497 100644 --- a/model_patch.py +++ b/model_patch.py @@ -13,31 +13,6 @@ def add_model_patch_option(model): return to -# Model needs current step number and cfg at inference step. It is possible to write a custom KSampler but I'd like to use ComfyUI's one. -# The first versions had modified_common_ksampler, but it broke custom KSampler nodes -def modified_sample(model, noise, positive, negative, cfg, device, sampler, sigmas, model_options={}, - latent_image=None, denoise_mask=None, callback=None, disable_pbar=False, seed=None): - ''' - Modified by BrushNet nodes - ''' - cfg_guider = comfy.samplers.CFGGuider(model) - cfg_guider.set_conds(positive, negative) - cfg_guider.set_cfg(cfg) - - ### Modified part ###################################################################### - # - to = add_model_patch_option(model) - to['model_patch']['all_sigmas'] = sigmas - - if math.isclose(cfg, 1.0) and model_options.get("disable_cfg1_optimization", False) == False: - to['model_patch']['free_guidance'] = False - else: - to['model_patch']['free_guidance'] = True - # - ####################################################################################### - - return cfg_guider.sample(noise, latent_image, sampler, sigmas, denoise_mask, callback, disable_pbar, seed) - # Patch model with model_function_wrapper def patch_model_function_wrapper(model, forward_patch): @@ -73,12 +48,12 @@ def patch_model_function_wrapper(model, forward_patch): t = model.model.model_sampling.timestep(timestep).float() # execute all patches for method in mp['forward']: - method(unet, xc, t, to) + method(unet, xc, t, to, control) return apply_model_method(x, timestep, **options_dict['c']) if "model_function_wrapper" in model.model_options and model.model_options["model_function_wrapper"]: - print('BrushNet is going to replace existing model_function_wrapper:', type(model.model_options["model_function_wrapper"])) + print('BrushNet is going to replace existing model_function_wrapper:', model.model_options["model_function_wrapper"]) model.set_model_unet_function_wrapper(brushnet_model_function_wrapper) to = add_model_patch_option(model) @@ -101,8 +76,56 @@ def patch_model_function_wrapper(model, forward_patch): mp['step'] = 0 mp['total_steps'] = 1 - # apply patch to code + # apply patches to code if comfy.samplers.sample.__doc__ is None or 'BrushNet' not in comfy.samplers.sample.__doc__: comfy.samplers.original_sample = comfy.samplers.sample comfy.samplers.sample = modified_sample + if comfy.ldm.modules.diffusionmodules.openaimodel.apply_control.__doc__ is None or \ + 'BrushNet' not in comfy.ldm.modules.diffusionmodules.openaimodel.apply_control.__doc__: + comfy.ldm.modules.diffusionmodules.openaimodel.original_apply_control = comfy.ldm.modules.diffusionmodules.openaimodel.apply_control + comfy.ldm.modules.diffusionmodules.openaimodel.apply_control = modified_apply_control + + +# Model needs current step number and cfg at inference step. It is possible to write a custom KSampler but I'd like to use ComfyUI's one. +# The first versions had modified_common_ksampler, but it broke custom KSampler nodes +def modified_sample(model, noise, positive, negative, cfg, device, sampler, sigmas, model_options={}, + latent_image=None, denoise_mask=None, callback=None, disable_pbar=False, seed=None): + ''' + Modified by BrushNet nodes + ''' + cfg_guider = comfy.samplers.CFGGuider(model) + cfg_guider.set_conds(positive, negative) + cfg_guider.set_cfg(cfg) + + ### Modified part ###################################################################### + # + to = add_model_patch_option(model) + to['model_patch']['all_sigmas'] = sigmas + + if math.isclose(cfg, 1.0) and model_options.get("disable_cfg1_optimization", False) == False: + to['model_patch']['free_guidance'] = False + else: + to['model_patch']['free_guidance'] = True + # + ####################################################################################### + + return cfg_guider.sample(noise, latent_image, sampler, sigmas, denoise_mask, callback, disable_pbar, seed) + + +# To use Controlnet with RAUNet it is much easier to modify apply_control a little +def modified_apply_control(h, control, name): + ''' + Modified by BrushNet nodes + ''' + if control is not None and name in control and len(control[name]) > 0: + ctrl = control[name].pop() + if ctrl is not None: + if h.shape[2] != ctrl.shape[2] or h.shape[3] != ctrl.shape[3]: + ctrl = torch.nn.functional.interpolate(ctrl, size=(h.shape[2], h.shape[3]), mode='bicubic').to(h.dtype).to(h.device) + try: + h += ctrl + except: + print.warning("warning control could not be applied {} {}".format(h.shape, ctrl.shape)) + return h + diff --git a/raunet_nodes.py b/raunet_nodes.py new file mode 100644 index 0000000..9d5494d --- /dev/null +++ b/raunet_nodes.py @@ -0,0 +1,158 @@ +import torch.nn.functional as F +import comfy + +from .model_patch import add_model_patch_option, patch_model_function_wrapper + + + +class RAUNet: + + @classmethod + def INPUT_TYPES(s): + return {"required": + { + "model": ("MODEL",), + "du_start": ("INT", {"default": 0, "min": 0, "max": 10000}), + "du_end": ("INT", {"default": 8, "min": 0, "max": 10000}), + "xa_start": ("INT", {"default": 8, "min": 0, "max": 10000}), + "xa_end": ("INT", {"default": 20, "min": 0, "max": 10000}), + }, + } + + CATEGORY = "inpaint" + RETURN_TYPES = ("MODEL",) + RETURN_NAMES = ("model",) + + FUNCTION = "model_update" + + def model_update(self, model, du_start, du_end, xa_start, xa_end): + + model = model.clone() + + add_raunet_patch(model, + du_start, + du_end, + xa_start, + xa_end) + + return (model,) + + +# This is main patch function +def add_raunet_patch(model, du_start, du_end, xa_start, xa_end): + + def raunet_forward(model, x, timesteps, transformer_options, control): + if 'model_patch' not in transformer_options: + print("RAUNet: 'model_patch' not in transformer_options, skip") + return + + mp = transformer_options['model_patch'] + is_SDXL = mp['SDXL'] + + if is_SDXL and type(model.input_blocks[6][0]) != comfy.ldm.modules.diffusionmodules.openaimodel.Downsample: + print('RAUNet: model is SDXL, but input[6] != Downsample, skip') + return + + if not is_SDXL and type(model.input_blocks[3][0]) != comfy.ldm.modules.diffusionmodules.openaimodel.Downsample: + print('RAUNet: model is not SDXL, but input[3] != Downsample, skip') + return + + if 'raunet' not in mp: + print('RAUNet: "raunet" not in model_patch options, skip') + return + + if is_SDXL: + block = model.input_blocks[6][0] + else: + block = model.input_blocks[3][0] + + total_steps = mp['total_steps'] + step = mp['step'] + + ro = mp['raunet'] + du_start = ro['du_start'] + du_end = ro['du_end'] + + if step >= du_start and step < du_end: + block.op.stride = (4, 4) + block.op.padding = (2, 2) + block.op.dilation = (2, 2) + else: + block.op.stride = (2, 2) + block.op.padding = (1, 1) + block.op.dilation = (1, 1) + + patch_model_function_wrapper(model, raunet_forward) + model.set_model_input_block_patch(in_xattn_patch) + model.set_model_output_block_patch(out_xattn_patch) + + to = add_model_patch_option(model) + mp = to['model_patch'] + if 'raunet' not in mp: + mp['raunet'] = {} + ro = mp['raunet'] + + ro['du_start'] = du_start + ro['du_end'] = du_end + ro['xa_start'] = xa_start + ro['xa_end'] = xa_end + + +def in_xattn_patch(h, transformer_options): + # both SDXL and SD15 = (input,4) + if transformer_options["block"] != ("input", 4): + # wrong block + return h + if 'model_patch' not in transformer_options: + print("RAUNet (i-x-p): 'model_patch' not in transformer_options") + return h + mp = transformer_options['model_patch'] + if 'raunet' not in mp: + print("RAUNet (i-x-p): 'raunet' not in model_patch options") + return h + + step = mp['step'] + ro = mp['raunet'] + xa_start = ro['xa_start'] + xa_end = ro['xa_end'] + + if step < xa_start or step >= xa_end: + return h + h = F.avg_pool2d(h, kernel_size=(2,2)) + return h + + +def out_xattn_patch(h, hsp, transformer_options): + if 'model_patch' not in transformer_options: + print("RAUNet (o-x-p): 'model_patch' not in transformer_options") + return h, hsp + mp = transformer_options['model_patch'] + if 'raunet' not in mp: + print("RAUNet (o-x-p): 'raunet' not in model_patch options") + return h + + step = mp['step'] + is_SDXL = mp['SDXL'] + ro = mp['raunet'] + xa_start = ro['xa_start'] + xa_end = ro['xa_end'] + + if is_SDXL: + if transformer_options["block"] != ("output", 5): + # wrong block + return h, hsp + else: + if transformer_options["block"] != ("output", 8): + # wrong block + return h, hsp + + if step < xa_start or step >= xa_end: + return h, hsp + #error in hidiffusion codebase, size * 2 for particular sizes only + #re_size = (int(h.shape[-2] * 2), int(h.shape[-1] * 2)) + re_size = (hsp.shape[-2], hsp.shape[-1]) + h = F.interpolate(h, size=re_size, mode='bicubic') + + return h, hsp + +