brushnet_nodes.py

import torch
import torchvision.transforms as T

import os
import folder_paths

import sys
# 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
import nodes
import latent_preview

from accelerate import init_empty_weights, load_checkpoint_and_dispatch

#import numpy as np
#import cv2
#from PIL import Image
#import yaml
#import importlib
#from contextlib import nullcontext
#from diffusers.loaders.single_file_utils import create_diffusers_vae_model_from_ldm
#from diffusers.image_processor import VaeImageProcessor

from .brushnet.brushnet import BrushNetModel

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__))
#original_config_file = os.path.join(current_directory, 'brushnet', 'v1-inference.yaml')
brushnet_config_file = os.path.join(current_directory, 'brushnet', 'brushnet.json')
brushnet_xl_config_file = os.path.join(current_directory, 'brushnet', 'brushnet_xl.json')

torch_dtype = torch.float16
sd15_scaling_factor = 0.18215
sdxl_scaling_factor = 0.13025

class BrushNetLoader:

    @classmethod
    def INPUT_TYPES(s):
        return {"required":
                    {    
                        "brushnet": (inpaint_safetensors(), ),
                     },
                }

    CATEGORY = "inpaint"
    RETURN_TYPES = ("BRMODEL",)
    RETURN_NAMES = ("brushnet",)

    FUNCTION = "brushnet_loading"

    def brushnet_loading(self, brushnet):
        brushnet_file = os.path.join(folder_paths.models_dir, "inpaint", brushnet)

        is_SDXL = False
        sd = comfy.utils.load_torch_file(brushnet_file)
        brushnet_down_block, brushnet_mid_block, brushnet_up_block = brushnet_blocks(sd)
        del sd
        if brushnet_down_block == 24 and brushnet_mid_block == 2 and brushnet_up_block == 30:
            print('Loading SD1.5 BrushNet model')
            is_SDXL = False
        elif brushnet_down_block == 18 and brushnet_mid_block == 2 and brushnet_up_block == 22:
            print('Loading SDXL BrushNet model')
            is_SDXL = True
        else:
            raise Exception("Unknown BrushNet model")

        with init_empty_weights():
            if is_SDXL:
                brushnet_config = BrushNetModel.load_config(brushnet_xl_config_file)
            else:
                brushnet_config = BrushNetModel.load_config(brushnet_config_file)
            brushnet_model = BrushNetModel.from_config(brushnet_config)

        print("BrushNet file", brushnet_file)

        brushnet_model = load_checkpoint_and_dispatch(
            brushnet_model,
            brushnet_file,
            device_map="auto",
            max_memory=None,
            offload_folder=None,
            offload_state_dict=False,
            dtype=torch_dtype,
            force_hooks=False,
        )

        print("BrushNet model is loaded, SDXL:", is_SDXL)

        return ({"brushnet": brushnet_model, "SDXL": is_SDXL},)

    
def inpaint_safetensors():
    inpaint_path = os.path.join(folder_paths.models_dir, 'inpaint')
    brushnet_path = os.path.join(inpaint_path, 'brushnet')
    abs_list = []
    for x in os.walk(inpaint_path):
        for name in x[2]:
            if 'safetensors' in name:
                abs_list.append(os.path.join(x[0], name))
    names = []
    for x in abs_list:
        remain = x
        y = ''
        while remain != inpaint_path:
            remain, folder = os.path.split(remain)
            if len(y) > 0:
                y = os.path.join(folder, y)
            else:
                y = folder
        names.append(y)     
    return names

def brushnet_blocks(sd):
    brushnet_down_block = 0
    brushnet_mid_block = 0
    brushnet_up_block = 0
    for key in sd:
        if 'brushnet_down_block' in key:
            brushnet_down_block += 1
        if 'brushnet_mid_block' in key:
            brushnet_mid_block += 1        
        if 'brushnet_up_block' in key:
            brushnet_up_block += 1
    return (brushnet_down_block, brushnet_mid_block, brushnet_up_block, )


class BrushNet:

    @classmethod
    def INPUT_TYPES(s):
        return {"required":
                    {    
                        "model": ("MODEL",),
                        "vae": ("VAE", ),
                        "image": ("IMAGE",),
                        "mask": ("MASK",),
                        "brushnet": ("BRMODEL", ),
                        "scale": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 10.0}),
                        "start_at": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 1.0}),
                        "end_at": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0}),
                     },
                }

    CATEGORY = "inpaint"
    RETURN_TYPES = ("MODEL","LATENT",)
    RETURN_NAMES = ("model","latent",)

    FUNCTION = "model_update"

    def model_update(self, model, vae, image, mask, brushnet, scale, start_at, end_at):

        is_SDXL = False
        if isinstance(model.model.model_config, comfy.supported_models.SD15):
            print('Base model type: SD1.5')
            is_SDXL = False
            if brushnet["SDXL"]:
                raise Exception("Base model is SD15, but BrushNet is SDXL type")    
        elif isinstance(model.model.model_config, comfy.supported_models.SDXL):
            print('Base model type: SDXL')
            is_SDXL = True
            if not brushnet["SDXL"]:
                raise Exception("Base model is SDXL, but BrushNet is SD15 type")    
            raise Exception("SDXL is not implemented yet")
        else:
            print('Base model type: ', type(model.model.model_config))
            raise Exception("Unsupported model type: " + str(type(model.model.model_config)))
        
        # prepare image and mask
        # no batches for original image and mask

        if image.shape[0] > 1:
            image = image[0][None,:,:,:]   
        if mask.shape[0] > 1:
            mask = mask[0][None,:,:]  

        width = image.shape[2]
        height = image.shape[1]

        print("BrushNet image,", image.shape)

        if mask.shape[2] != width or mask.shape[1] != height:
            raise Exception("Image and mask should be the same size")

        masked_image = image * (1.0 - mask[:,:,:,None])

        if hasattr(model.model.model_config, 'latent_format') and hasattr(model.model.model_config.latent_format, 'scale_factor'):
            scaling_factor = model.model.model_config.latent_format.scale_factor
        elif is_SDXL:
            scaling_factor = sdxl_scaling_factor
        else:
            scaling_factor = sd15_scaling_factor

        processed_image = torch.cat([masked_image] * 2).to(vae.device)
        image_latents = vae.encode(processed_image[:,:,:,:3]) * scaling_factor

        processed_mask = torch.cat([1. - mask[None,:,:,:]] * 2)
        interpolated_mask = torch.nn.functional.interpolate(
                    processed_mask, 
                    size=(
                        image_latents.shape[-2], 
                        image_latents.shape[-1]
                    )
                )
        interpolated_mask = interpolated_mask.to(image_latents.device)

        print('BrushNet: image latents shape =', image_latents.shape, 'interpolated mask shape =', interpolated_mask.shape)

        conditioning_latents = torch.concat([image_latents, interpolated_mask], 1).to(dtype=torch_dtype).to(brushnet['brushnet'].device)

        # apply patches to code

        if nodes.common_ksampler.__doc__ is None or 'BrushNet' not in nodes.common_ksampler.__doc__:
            nodes.original_common_ksampler = nodes.common_ksampler
            nodes.common_ksampler = modified_common_ksampler

        if comfy.ldm.modules.diffusionmodules.openaimodel.forward_timestep_embed.__doc__ is None or \
            'BrushNet' not in comfy.ldm.modules.diffusionmodules.openaimodel.forward_timestep_embed.__doc__:
            comfy.ldm.modules.diffusionmodules.openaimodel.original_forward_timestep_embed =  \
            comfy.ldm.modules.diffusionmodules.openaimodel.forward_timestep_embed
            comfy.ldm.modules.diffusionmodules.openaimodel.forward_timestep_embed = modified_forward_timestep_embed

        # apply patch to model

        brushnet_conditioning_scale = scale
        control_guidance_start = start_at
        control_guidance_end = end_at

        add_brushnet_patch(model, brushnet["brushnet"], conditioning_latents, 
                           [brushnet_conditioning_scale, control_guidance_start, control_guidance_end])
        
        if is_SDXL:
            add_model_patch(model, brushnet_inference, ('input', 0), (0, 'before'))
            input_blocks = [[0,0],[1,0],[2,0],[3,0],[4,1],[5,1],[6,0],[7,1],[8,1]]
            for i, j in input_blocks:
                add_model_patch(model, apply_brushnet, ('input', i), (j, 'after'))
            add_model_patch(model, apply_brushnet, ('middle', 0), (2, 'after'))
            output_blocks = [[0,1],[1,1],[2,1],[2,2],[3,1],[4,1],[5,1],[5,2],[6,0],[7,0],[8,0]]
            for i, j in output_blocks:
                add_model_patch(model, apply_brushnet, ('output', i), (j, 'after'))
        else:
            add_model_patch(model, brushnet_inference, ('input', 0), (0, 'before'))
            input_blocks = [[0,0],[1,1],[2,1],[3,0],[4,1],[5,1],[6,0],[7,1],[8,1],[9,0],[10,0],[11,0]]
            for i, j in input_blocks:
                add_model_patch(model, apply_brushnet, ('input', i), (j, 'after'))
            add_model_patch(model, apply_brushnet, ('middle', 0), (2, 'after'))
            output_blocks = [[0,0],[1,0],[2,0],[2,1],[3,1],[4,1],[5,1],[5,2],[6,1],[7,1],[8,1],[8,2],[9,1],[10,1],[11,1]]
            for i, j in output_blocks:
                add_model_patch(model, apply_brushnet, ('output', i), (j, 'after'))

        latent = torch.zeros([1, 4, conditioning_latents.shape[2], conditioning_latents.shape[3]], device=brushnet['brushnet'].device)

        return (model, {"samples":latent},)


class BlendInpaint:

    @classmethod
    def INPUT_TYPES(s):
        return {"required":
                    {    
                        "inpaint": ("IMAGE",),
                        "original": ("IMAGE",),
                        "mask": ("MASK",),
                        "kernel": ("INT", {"default": 10, "min": 1, "max": 1000}),
                        "sigma": ("FLOAT", {"default": 10.0, "min": 0.01, "max": 1000}),
                     },
                }

    CATEGORY = "inpaint"
    RETURN_TYPES = ("IMAGE","MASK",)
    RETURN_NAMES = ("image","MASK",)

    FUNCTION = "blend_inpaint"

    def blend_inpaint(self, inpaint: torch.Tensor, original: torch.Tensor, mask: torch.Tensor, kernel: int, sigma:int) -> Tuple[torch.Tensor]:

        # no batches over mask and original image
        if len(mask.shape) > 2:
            mask = mask[0]
        if len(original.shape) > 3:
            original = original[0]

        if kernel % 2 == 0:
            kernel += 1
        transform = T.GaussianBlur(kernel_size=(kernel, kernel), sigma=(sigma, sigma))
        blurred_mask = transform(mask[None,None,:,:]).to(original.device).to(original.dtype)

        ret = []
        for result in inpaint:
            ret.append(original * (1.0 - blurred_mask[0][0][:,:,None]) + result.to(original.device) * blurred_mask[0][0][:,:,None])

        return (torch.stack(ret), blurred_mask[0],)

'''
import pip

class TestNode:

    @classmethod
    def INPUT_TYPES(s):
        return {"required":
                    {    
                        "image": ("IMAGE",),
                     },
                }

    CATEGORY = "inpaint"
    RETURN_TYPES = ()
    RETURN_NAMES = ()

    FUNCTION = "test_node"

    def test_node(self):
        print('Installed packages')
        print(pip.utils.get_installed_distributions())
'''



# Unfortunately, ModelPatcher does not have necessary hooks to patch, so we have to patch the code instead
def modified_forward_timestep_embed(block, x, emb, context=None, transformer_options={}, 
                           output_shape=None, time_context=None, 
                           num_video_frames=None, image_only_indicator=None):
    '''
    Modified by BrushNet nodes
    '''

    if 'model_patch' not in transformer_options or transformer_options['block'] not in transformer_options['model_patch']:
        return comfy.ldm.modules.diffusionmodules.openaimodel.original_forward_timestep_embed(block, x, emb, context, 
                                                                                              transformer_options,
                                                                                              output_shape, time_context, 
                                                                                              num_video_frames, image_only_indicator)
        
    block_patch = transformer_options['model_patch'][transformer_options['block']]

    for i, layer in enumerate(block):

        if (i, 'before') in block_patch:
            x = block_patch[(i, 'before')](x, emb, context, (transformer_options['block'], i, 'before'), transformer_options)

        if isinstance(layer, comfy.ldm.modules.diffusionmodules.openaimodel.VideoResBlock):
            x = layer(x, emb, num_video_frames, image_only_indicator)
        elif isinstance(layer, comfy.ldm.modules.diffusionmodules.openaimodel.TimestepBlock):
            x = layer(x, emb)
        elif isinstance(layer, comfy.ldm.modules.attention.SpatialVideoTransformer):
            x = layer(x, context, time_context, num_video_frames, image_only_indicator, transformer_options)
            if "transformer_index" in transformer_options:
                transformer_options["transformer_index"] += 1
        elif isinstance(layer, comfy.ldm.modules.attention.SpatialTransformer):
            x = layer(x, context, transformer_options)
            if "transformer_index" in transformer_options:
                transformer_options["transformer_index"] += 1
        elif isinstance(layer, comfy.ldm.modules.diffusionmodules.openaimodel.Upsample):
            x = layer(x, output_shape=output_shape)
        else:
            x = layer(x)

        if (i, 'after') in block_patch:
            x = block_patch[(i, 'after')](x, emb, context, (transformer_options['block'], i, 'after'), transformer_options)
        
    return x

# Model needs current step number at inference stepю. It is possible to write a custom KSampler but we'd like to use ComfyUI's one.
def modified_common_ksampler(model, seed, steps, cfg, sampler_name, scheduler, positive, negative, latent, denoise=1.0, 
                             disable_noise=False, start_step=None, last_step=None, force_full_denoise=False):
    '''
    Modified by BrushNet nodes
    '''
    latent_image = latent["samples"]
    if disable_noise:
        noise = torch.zeros(latent_image.size(), dtype=latent_image.dtype, layout=latent_image.layout, device="cpu")
    else:
        batch_inds = latent["batch_index"] if "batch_index" in latent else None
        noise = comfy.sample.prepare_noise(latent_image, seed, batch_inds)

    noise_mask = None
    if "noise_mask" in latent:
        noise_mask = latent["noise_mask"]

    #######################################################################################
    #
    latent_preview_callback = latent_preview.prepare_callback(model, steps)

    to = add_model_patch_option(model)
    to['model_patch']['step'] = 0
    to['model_patch']['total_steps'] = steps

    def callback(step, x0, x, total_steps):
        to['model_patch']['step'] = step + 1
        latent_preview_callback(steps, x0, x, total_steps)
    #
    #######################################################################################
    
    disable_pbar = not comfy.utils.PROGRESS_BAR_ENABLED
    samples = comfy.sample.sample(model, noise, steps, cfg, sampler_name, scheduler, positive, negative, latent_image,
                                  denoise=denoise, disable_noise=disable_noise, start_step=start_step, last_step=last_step,
                                  force_full_denoise=force_full_denoise, noise_mask=noise_mask, callback=callback, 
                                  disable_pbar=disable_pbar, seed=seed)
    out = latent.copy()
    out["samples"] = samples
    return (out, )

def add_model_patch_option(model):
    if 'transformer_options' not in model.model_options:
        model.model_options['transformer_options'] = {}
    to = model.model_options['transformer_options']
    if "model_patch" not in to:
        to["model_patch"] = {}
    return to

def add_model_patch(model, patch, block, key, replace=True):
    to = add_model_patch_option(model)
    if block not in to["model_patch"]:
        to["model_patch"][block] = {}
    if key not in to["model_patch"][block]:
        to["model_patch"][block][key] = []

    if replace:
        to["model_patch"][block][key] = patch
    else:
        to["model_patch"][block][key].append(patch)

def add_brushnet_patch(model, brushnet, conditioning_latents, controls):
    to = add_model_patch_option(model)
    to['model_patch']['brushnet_model'] = brushnet
    to['model_patch']['brushnet_latents'] = conditioning_latents
    to['model_patch']['brushnet_controls'] = controls
    to['model_patch']['input_samples'] = []
    to['model_patch']['middle_sample'] = 0
    to['model_patch']['output_samples'] = []

def brushnet_inference(x, emb, context, loc, transformer_options):
    # x : sample
    # emb : time embedding
    # context[0] == neg_cond[0][0][0]
    # context[1] == pos_cond[0][0][0]
    # loc = (('input'|'middle'|'output', i), j, 'before'|'after')

    if 'model_patch' not in transformer_options:
        raise Exception('BrushNet inference: there is no model_patch in transformer_options')
    
    mp = transformer_options['model_patch']
    brushnet = mp['brushnet_model']
    if isinstance(brushnet, BrushNetModel):
        conditioning_latents = mp['brushnet_latents']
        step = mp['step']
        total_steps = mp['total_steps']
        brushnet_conditioning_scale, control_guidance_start, control_guidance_end = mp['brushnet_controls']
        brushnet_keep = []
        for i in range(total_steps):
            keeps = [
                1.0 - float(i / total_steps < s or (i + 1) / total_steps > e)
                for s, e in zip([control_guidance_start], [control_guidance_end])
            ]
            brushnet_keep.append(keeps[0])
        cond_scale = brushnet_conditioning_scale * brushnet_keep[step]

        down_samples, mid_sample, up_samples = brushnet(x,
                                                        encoder_hidden_states=context[1],
                                                        brushnet_cond=conditioning_latents,
                                                        time_emb=emb,
                                                        conditioning_scale=cond_scale,
                                                        guess_mode=False,
                                                        return_dict=False,)
        transformer_options['model_patch']['input_samples'] = down_samples
        transformer_options['model_patch']['middle_sample'] = mid_sample
        transformer_options['model_patch']['output_samples'] = up_samples
    else:
        raise Exception('BrushNet model is not a BrushNetModel class')
        
    return x

def apply_brushnet(x, emb, context, loc, transformer_options):
    if loc[0][0] == 'input':
        if len(transformer_options['model_patch']['input_samples']) > 0:
            return x + transformer_options['model_patch']['input_samples'].pop(0)
        else:
            raise Exception('BrushNet: something is not right, input samples are empty, ' + str(loc))
        
    elif loc[0][0] == 'middle':
        return x + transformer_options['model_patch']['middle_sample']
    else:
        if len(transformer_options['model_patch']['output_samples']) > 0:
            return x + transformer_options['model_patch']['output_samples'].pop(0)
        else:
            raise Exception('BrushNet: something is not right, output samples are empty, ' + str(loc))

#def image_from_tensor(t: torch.Tensor) -> np.ndarray:
#    image_np = t.numpy() 
#    # Convert the numpy array back to the original range (0-255) and data type (uint8)
#    image_np = (image_np * 255).astype(np.uint8)
#    return image_np
#
#def numpy_to_tensor(array: np.ndarray) -> torch.Tensor:
#    """Convert a numpy array to a tensor and scale its values from 0-255 to 0-1."""
#    array = array.astype(np.float32) / 255.0
#    return torch.from_numpy(array)[None,]