(model, forward_patch, remove=False)
| 14 | |
| 15 | # Patch model with model_function_wrapper |
| 16 | def 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 = None |
| 22 | 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 execute |
| 29 | 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 | all_sigmas = mp['all_sigmas'] |
| 36 | sigma = to['sigmas'][0].item() |
| 37 | total_steps = all_sigmas.shape[0] - 1 |
| 38 | step = torch.argmin((all_sigmas - sigma).abs()).item() |
| 39 | |
| 40 | mp['step'] = step |
| 41 | mp['total_steps'] = total_steps |
| 42 | |
| 43 | # comfy.model_base.apply_model |
| 44 | xc = model.model.model_sampling.calculate_input(timestep, x) |
| 45 | if 'c_concat' in options_dict['c'] and options_dict['c']['c_concat'] is not None: |
| 46 | xc = torch.cat([xc] + [options_dict['c']['c_concat']], dim=1) |
| 47 | t = model.model.model_sampling.timestep(timestep).float() |
| 48 | # execute all patches |
| 49 | for method in mp['forward']: |
| 50 | method(unet, xc, t, to, control) |
| 51 | |
| 52 | return apply_model_method(x, timestep, **options_dict['c']) |
| 53 | |
| 54 | if "model_function_wrapper" in model.model_options and model.model_options["model_function_wrapper"]: |
| 55 | print('BrushNet is going to replace existing model_function_wrapper:', model.model_options["model_function_wrapper"]) |
| 56 | model.set_model_unet_function_wrapper(brushnet_model_function_wrapper) |
| 57 | |
| 58 | to = add_model_patch_option(model) |
| 59 | mp = to['model_patch'] |
| 60 | |
| 61 | if isinstance(model.model.model_config, comfy.supported_models.SD15): |
| 62 | mp['SDXL'] = False |
| 63 | elif isinstance(model.model.model_config, comfy.supported_models.SDXL): |
| 64 | mp['SDXL'] = True |
| 65 | else: |
| 66 | print('Base model type: ', type(model.model.model_config)) |
| 67 | raise Exception("Unsupported model type: ", type(model.model.model_config)) |
| 68 | |
| 69 | if 'forward' not in mp: |
| 70 | mp['forward'] = [] |
| 71 | |
| 72 | if remove: |
| 73 | if forward_patch in mp['forward']: |
no test coverage detected