From afe771460159e835e50d53506f9b8ecfaa489cb8 Mon Sep 17 00:00:00 2001 From: nullquant Date: Tue, 7 May 2024 18:00:20 +0300 Subject: [PATCH] dtype switch --- README.md | 11 ++++++- brushnet_nodes.py | 78 +++++++++++++++++++++++++++++++---------------- 2 files changed, 61 insertions(+), 28 deletions(-) diff --git a/README.md b/README.md index 6d710d2..a3182ea 100644 --- a/README.md +++ b/README.md @@ -20,8 +20,8 @@ Apr 11, 2024. Initial commit. - [x] BrushNet SDXL - [x] PowerPaint v2 -- [ ] Compatibility with `jank HiDiffusion` and similar nodes - [ ] Image batch +- [ ] Compatibility with `jank HiDiffusion` and similar nodes ## Installation @@ -59,9 +59,18 @@ Below is an example for the intended workflow. The [workflow](example/BrushNet_b ### Parameters +#### Brushnet Loader + +- `dtype`, defaults to `torch.float16`. The torch.dtype of BrushNet. If you have old GPU or NVIDIA 16 series card try to switch to `torch.float32`. + +#### Brushnet + - `scale`, defaults to 1.0: The "strength" of BrushNet. The outputs of the BrushNet are multiplied by `scale` before they are added to the residual in the original unet. - `start_at`, defaults to 0: step at which the BrushNet starts applying. - `end_at`, defaults to 10000: step at which the BrushNet stops applying. + +#### PowerPaint + - `CLIP`: PowerPaint CLIP that should be passed from PowerPaintCLIPLoader node. - `fitting`: PowerPaint fitting degree. - `function`: PowerPaint function, see its [page](https://github.com/open-mmlab/PowerPaint) for details. diff --git a/brushnet_nodes.py b/brushnet_nodes.py index 0ae8ea0..9b5dc82 100644 --- a/brushnet_nodes.py +++ b/brushnet_nodes.py @@ -5,6 +5,7 @@ import os import folder_paths 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) @@ -30,10 +31,10 @@ brushnet_config_file = os.path.join(current_directory, 'brushnet', 'brushnet.jso brushnet_xl_config_file = os.path.join(current_directory, 'brushnet', 'brushnet_xl.json') powerpaint_config_file = os.path.join(current_directory,'brushnet', 'powerpaint.json') -torch_dtype = torch.float16 sd15_scaling_factor = 0.18215 sdxl_scaling_factor = 0.13025 + class BrushNetLoader: @classmethod @@ -41,6 +42,7 @@ class BrushNetLoader: return {"required": { "brushnet": (get_files_with_extension('inpaint'), ), + "dtype": (['float16', 'bfloat16', 'float32', 'float64'], ), }, } @@ -50,7 +52,7 @@ class BrushNetLoader: FUNCTION = "brushnet_loading" - def brushnet_loading(self, brushnet): + def brushnet_loading(self, brushnet, dtype): brushnet_file = os.path.join(folder_paths.models_dir, "inpaint", brushnet) is_SDXL = False @@ -89,6 +91,15 @@ class BrushNetLoader: else: print("BrushNet model file:", brushnet_file) + if dtype == 'float16': + torch_dtype = torch.float16 + elif dtype == 'bfloat16': + torch_dtype = torch.bfloat16 + elif dtype == 'float32': + torch_dtype = torch.float32 + else: + torch_dtype = torch.float64 + brushnet_model = load_checkpoint_and_dispatch( brushnet_model, brushnet_file, @@ -107,7 +118,7 @@ class BrushNetLoader: else: print("BrushNet SD1.5 model is loaded") - return ({"brushnet": brushnet_model, "SDXL": is_SDXL, "PP": is_PP},) + return ({"brushnet": brushnet_model, "SDXL": is_SDXL, "PP": is_PP, "dtype": torch_dtype}, ) class PowerPaintCLIPLoader: @@ -205,6 +216,8 @@ class PowerPaint: else: scaling_factor = sd15_scaling_factor + torch_dtype = powerpaint['dtype'] + # prepare conditioning latents conditioning_latents, conditioning_latents2 = get_image_latents(masked_image, mask, vae, scaling_factor) conditioning_latents = conditioning_latents.to(dtype=torch_dtype).to(powerpaint['brushnet'].device) @@ -275,6 +288,7 @@ class PowerPaint: add_brushnet_patch(model, powerpaint['brushnet'], + torch_dtype, conditioning_latents, conditioning_latents2, (brushnet_conditioning_scale, control_guidance_start, control_guidance_end), @@ -336,6 +350,8 @@ class BrushNet: else: scaling_factor = sd15_scaling_factor + torch_dtype = brushnet['dtype'] + # prepare conditioning latents conditioning_latents, conditioning_latents2 = get_image_latents(masked_image, mask, vae, scaling_factor) conditioning_latents = conditioning_latents.to(dtype=torch_dtype).to(brushnet['brushnet'].device) @@ -383,6 +399,7 @@ class BrushNet: add_brushnet_patch(model, brushnet['brushnet'], + torch_dtype, conditioning_latents, conditioning_latents2, (brushnet_conditioning_scale, control_guidance_start, control_guidance_end), @@ -480,6 +497,7 @@ def brushnet_blocks(sd): brushnet_up_block += 1 return (brushnet_down_block, brushnet_mid_block, brushnet_up_block, len(sd)) + # Check models compatibility def check_compatibilty(model, brushnet): is_SDXL = False @@ -536,36 +554,36 @@ def prepare_image(image, mask): # Prepare conditioning_latents +@torch.inference_mode() def get_image_latents(masked_image, mask, vae, scaling_factor): - with torch.no_grad(): - processed_image = masked_image.to(vae.device) - processed_image2 = torch.cat([masked_image] * 2).to(vae.device) + processed_image = masked_image.to(vae.device) + processed_image2 = torch.cat([masked_image] * 2).to(vae.device) - image_latents = vae.encode(processed_image[:,:,:,:3]) * scaling_factor - image_latents2 = vae.encode(processed_image2[:,:,:,:3]) * scaling_factor + image_latents = vae.encode(processed_image[:,:,:,:3]) * scaling_factor + image_latents2 = vae.encode(processed_image2[:,:,:,:3]) * scaling_factor - processed_mask = 1. - mask[None,:,:,:] - interpolated_mask = torch.nn.functional.interpolate( - processed_mask, - size=( - image_latents.shape[-2], - image_latents.shape[-1] - ) + processed_mask = 1. - mask[None,:,:,:] + 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) + ) + interpolated_mask = interpolated_mask.to(image_latents.device) - processed_mask2 = torch.cat([1. - mask[None,:,:,:]] * 2) - interpolated_mask2 = torch.nn.functional.interpolate( - processed_mask2, - size=( - image_latents2.shape[-2], - image_latents2.shape[-1] - ) + processed_mask2 = torch.cat([1. - mask[None,:,:,:]] * 2) + interpolated_mask2 = torch.nn.functional.interpolate( + processed_mask2, + size=( + image_latents2.shape[-2], + image_latents2.shape[-1] ) - interpolated_mask2 = interpolated_mask2.to(image_latents2.device) + ) + interpolated_mask2 = interpolated_mask2.to(image_latents2.device) - conditioning_latents = torch.concat([image_latents, interpolated_mask], 1) - conditioning_latents2 = torch.concat([image_latents2, interpolated_mask2], 1) + conditioning_latents = torch.concat([image_latents, interpolated_mask], 1) + conditioning_latents2 = torch.concat([image_latents2, interpolated_mask2], 1) print('BrushNet CL: image_latents shape =', image_latents2.shape, 'interpolated_mask shape =', interpolated_mask2.shape) @@ -634,6 +652,7 @@ def modified_sample(model, noise, positive, negative, cfg, device, sampler, sigm return cfg_guider.sample(noise, latent_image, sampler, sigmas, denoise_mask, callback, disable_pbar, seed) # Main function where magic happens +@torch.inference_mode() def brushnet_inference(x, timesteps, transformer_options): if 'model_patch' not in transformer_options: print('BrushNet inference: there is no model_patch in transformer_options') @@ -644,6 +663,7 @@ def brushnet_inference(x, timesteps, transformer_options): return ([], 0, []) brushnet = mp['brushnet_model'] if isinstance(brushnet, BrushNetModel) or isinstance(brushnet, PowerPaintModel): + torch_dtype = mp['brushnet_dtype'] cl = mp['brushnet_latents'] cl2 = mp['brushnet_latents2'] brushnet_conditioning_scale, control_guidance_start, control_guidance_end = mp['brushnet_controls'] @@ -752,6 +772,9 @@ def brushnet_inference(x, timesteps, transformer_options): added_cond_kwargs=added_cond_kwargs, return_dict=False, ) + + + else: print('BrushNet model is not a BrushNetModel class') return ([], 0, []) @@ -766,12 +789,13 @@ def add_model_patch_option(model): return to # This is main patch function -def add_brushnet_patch(model, brushnet, conditioning_latents, conditioning_latents2, +def add_brushnet_patch(model, brushnet, torch_dtype, conditioning_latents, conditioning_latents2, controls, prompt_embeds, add_text_embeds, add_time_ids, prompt_embeds2, add_text_embeds2, add_time_ids2): to = add_model_patch_option(model) to['model_patch']['brushnet_model'] = brushnet + to['model_patch']['brushnet_dtype'] = torch_dtype to['model_patch']['brushnet_latents'] = conditioning_latents to['model_patch']['brushnet_latents2'] = conditioning_latents2 to['model_patch']['brushnet_controls'] = controls