erikayurika/brush

1import torch2import comfy3 4 5# Check and add 'model_patch' to model.model_options['transformer_options']6def add_model_patch_option(model):7    if 'transformer_options' not in model.model_options:8        model.model_options['transformer_options'] = {}9    to = model.model_options['transformer_options']10    if "model_patch" not in to:11        to["model_patch"] = {}12    return to13 14 15# Patch model with model_function_wrapper16def patch_model_function_wrapper(model, forward_patch, remove=False):17 18    def brushnet_model_function_wrapper(apply_model_method, options_dict):19        to = options_dict['c']['transformer_options']20 21        control = None22        if 'control' in options_dict['c']:23            control = options_dict['c']['control']24        25        x = options_dict['input']26        timestep = options_dict['timestep']27        28        # check if there are patches to execute29        if 'model_patch' not in to or 'forward' not in to['model_patch']:30            return apply_model_method(x, timestep, **options_dict['c'])31 32        mp = to['model_patch']33        unet = mp['unet']34        35 36 37        #print(model.get_model_object("model_sampling").sigmas, len(model.get_model_object("model_sampling").sigmas))38        #print(mp['all_sigmas'], len(mp['all_sigmas']))39 40 41        all_sigmas = mp['all_sigmas']42        sigma = to['sigmas'][0].item()43        total_steps = all_sigmas.shape[0] - 144        step = torch.argmin((all_sigmas - sigma).abs()).item()45 46        mp['step'] = step47        mp['total_steps'] = total_steps48 49        # comfy.model_base.apply_model50        xc = model.model.model_sampling.calculate_input(timestep, x)51        if 'c_concat' in options_dict['c'] and options_dict['c']['c_concat'] is not None:52            xc = torch.cat([xc] + [options_dict['c']['c_concat']], dim=1)53        t = model.model.model_sampling.timestep(timestep).float()54        # execute all patches        55        for method in mp['forward']:56            method(unet, xc, t, to, control)57            58        return apply_model_method(x, timestep, **options_dict['c'])59    60    if "model_function_wrapper" in model.model_options and model.model_options["model_function_wrapper"]:61        print('BrushNet is going to replace existing model_function_wrapper:', model.model_options["model_function_wrapper"])62    model.set_model_unet_function_wrapper(brushnet_model_function_wrapper)   63 64    to = add_model_patch_option(model)65    mp = to['model_patch']66 67    if isinstance(model.model.model_config, comfy.supported_models.SD15):68        mp['SDXL'] = False69    elif isinstance(model.model.model_config, comfy.supported_models.SDXL):70        mp['SDXL'] = True71    else:72        print('Base model type: ', type(model.model.model_config))73        raise Exception("Unsupported model type: ", type(model.model.model_config))74 75    if 'forward' not in mp:76        mp['forward'] = []77 78    if remove:79        if forward_patch in mp['forward']:80            mp['forward'].remove(forward_patch)81    else:82        mp['forward'].append(forward_patch)83 84    mp['unet'] = model.model.diffusion_model85    mp['step'] = 086    mp['total_steps'] = 187 88    # apply patches to code89    if comfy.samplers.sample.__doc__ is None or 'BrushNet' not in comfy.samplers.sample.__doc__:90        comfy.samplers.original_sample = comfy.samplers.sample91        comfy.samplers.sample = modified_sample92 93    if comfy.ldm.modules.diffusionmodules.openaimodel.apply_control.__doc__ is None or \94        'BrushNet' not in comfy.ldm.modules.diffusionmodules.openaimodel.apply_control.__doc__:95        comfy.ldm.modules.diffusionmodules.openaimodel.original_apply_control = comfy.ldm.modules.diffusionmodules.openaimodel.apply_control96        comfy.ldm.modules.diffusionmodules.openaimodel.apply_control = modified_apply_control97 98 99# 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.100# The first versions had modified_common_ksampler, but it broke custom KSampler nodes101def modified_sample(model, noise, positive, negative, cfg, device, sampler, sigmas, model_options={}, 102           latent_image=None, denoise_mask=None, callback=None, disable_pbar=False, seed=None):103    '''104    Modified by BrushNet nodes105    '''106    cfg_guider = comfy.samplers.CFGGuider(model)107    cfg_guider.set_conds(positive, negative)108    cfg_guider.set_cfg(cfg)109 110    ### Modified part ######################################################################111    #112    to = add_model_patch_option(model)113    to['model_patch']['all_sigmas'] = sigmas114    #115    #sigma_start = model.get_model_object("model_sampling").percent_to_sigma(start_at)116    #sigma_end = model.get_model_object("model_sampling").percent_to_sigma(end_at)117    #118    #119    #if math.isclose(cfg, 1.0) and model_options.get("disable_cfg1_optimization", False) == False:120    #    to['model_patch']['free_guidance'] = False121    #else:122    #    to['model_patch']['free_guidance'] = True123    #124    #######################################################################################125       126    return cfg_guider.sample(noise, latent_image, sampler, sigmas, denoise_mask, callback, disable_pbar, seed)127 128 129# To use Controlnet with RAUNet it is much easier to modify apply_control a little130def modified_apply_control(h, control, name):131    '''132    Modified by BrushNet nodes133    '''134    if control is not None and name in control and len(control[name]) > 0:135        ctrl = control[name].pop()136        if ctrl is not None:137            if h.shape[2] != ctrl.shape[2] or h.shape[3] != ctrl.shape[3]:138                ctrl = torch.nn.functional.interpolate(ctrl, size=(h.shape[2], h.shape[3]), mode='bicubic').to(h.dtype).to(h.device)                    139            try:140                h += ctrl141            except:142                print.warning("warning control could not be applied {} {}".format(h.shape, ctrl.shape))143    return h    144 145