RAUNet
This commit is contained in:
nullquant
@@ -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:
|
||||
|
||||
|
||||
|
||||
[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.
|
||||
|
||||
|
||||
|
||||
[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.
|
||||
@@ -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)
|
||||
|
||||
@@ -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",
|
||||
}
|
||||
|
||||
+17
-22
@@ -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
|
||||
|
||||
Binary file not shown.
|
After Width: | Height: | Size: 1.7 MiB |
Binary file not shown.
|
After Width: | Height: | Size: 1.8 MiB |
File diff suppressed because one or more lines are too long
File diff suppressed because one or more lines are too long
+51
-28
@@ -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
|
||||
|
||||
|
||||
+158
@@ -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
|
||||
|
||||
|
||||
Reference in New Issue
Block a user