From 94a5a67c32ded77efbd1b47e8d0257ddab9ac3bc Mon Sep 17 00:00:00 2001 From: comfyanonymous Date: Fri, 29 Mar 2024 14:43:24 -0400 Subject: [PATCH 01/29] Cleanup to support different types of inpaint models. --- comfy/model_base.py | 43 ++++++++++++++++++++++--------------------- 1 file changed, 22 insertions(+), 21 deletions(-) diff --git a/comfy/model_base.py b/comfy/model_base.py index bc019de53..6f530d2fa 100644 --- a/comfy/model_base.py +++ b/comfy/model_base.py @@ -66,7 +66,8 @@ class BaseModel(torch.nn.Module): self.adm_channels = unet_config.get("adm_in_channels", None) if self.adm_channels is None: self.adm_channels = 0 - self.inpaint_model = False + + self.concat_keys = () logging.info("model_type {}".format(model_type.name)) logging.debug("adm {}".format(self.adm_channels)) @@ -107,8 +108,7 @@ class BaseModel(torch.nn.Module): def extra_conds(self, **kwargs): out = {} - if self.inpaint_model: - concat_keys = ("mask", "masked_image") + if len(self.concat_keys) > 0: cond_concat = [] denoise_mask = kwargs.get("concat_mask", kwargs.get("denoise_mask", None)) concat_latent_image = kwargs.get("concat_latent_image", None) @@ -125,24 +125,16 @@ class BaseModel(torch.nn.Module): concat_latent_image = utils.resize_to_batch_size(concat_latent_image, noise.shape[0]) - if len(denoise_mask.shape) == len(noise.shape): - denoise_mask = denoise_mask[:,:1] + if denoise_mask is not None: + if len(denoise_mask.shape) == len(noise.shape): + denoise_mask = denoise_mask[:,:1] - denoise_mask = denoise_mask.reshape((-1, 1, denoise_mask.shape[-2], denoise_mask.shape[-1])) - if denoise_mask.shape[-2:] != noise.shape[-2:]: - denoise_mask = utils.common_upscale(denoise_mask, noise.shape[-1], noise.shape[-2], "bilinear", "center") - denoise_mask = utils.resize_to_batch_size(denoise_mask.round(), noise.shape[0]) + denoise_mask = denoise_mask.reshape((-1, 1, denoise_mask.shape[-2], denoise_mask.shape[-1])) + if denoise_mask.shape[-2:] != noise.shape[-2:]: + denoise_mask = utils.common_upscale(denoise_mask, noise.shape[-1], noise.shape[-2], "bilinear", "center") + denoise_mask = utils.resize_to_batch_size(denoise_mask.round(), noise.shape[0]) - def blank_inpaint_image_like(latent_image): - blank_image = torch.ones_like(latent_image) - # these are the values for "zero" in pixel space translated to latent space - blank_image[:,0] *= 0.8223 - blank_image[:,1] *= -0.6876 - blank_image[:,2] *= 0.6364 - blank_image[:,3] *= 0.1380 - return blank_image - - for ck in concat_keys: + for ck in self.concat_keys: if denoise_mask is not None: if ck == "mask": cond_concat.append(denoise_mask.to(device)) @@ -152,7 +144,7 @@ class BaseModel(torch.nn.Module): if ck == "mask": cond_concat.append(torch.ones_like(noise)[:,:1]) elif ck == "masked_image": - cond_concat.append(blank_inpaint_image_like(noise)) + cond_concat.append(self.blank_inpaint_image_like(noise)) data = torch.cat(cond_concat, dim=1) out['c_concat'] = comfy.conds.CONDNoiseShape(data) @@ -221,7 +213,16 @@ class BaseModel(torch.nn.Module): return unet_state_dict def set_inpaint(self): - self.inpaint_model = True + self.concat_keys = ("mask", "masked_image") + def blank_inpaint_image_like(latent_image): + blank_image = torch.ones_like(latent_image) + # these are the values for "zero" in pixel space translated to latent space + blank_image[:,0] *= 0.8223 + blank_image[:,1] *= -0.6876 + blank_image[:,2] *= 0.6364 + blank_image[:,3] *= 0.1380 + return blank_image + self.blank_inpaint_image_like = blank_inpaint_image_like def memory_required(self, input_shape): if comfy.model_management.xformers_enabled() or comfy.model_management.pytorch_attention_flash_attention(): From 96b4c757cf55698d50c782cf0f9462bc40ec4c19 Mon Sep 17 00:00:00 2001 From: comfyanonymous Date: Sat, 30 Mar 2024 11:52:11 -0400 Subject: [PATCH 02/29] Add log to debug custom nodes that hang when imported. --- nodes.py | 1 + 1 file changed, 1 insertion(+) diff --git a/nodes.py b/nodes.py index 6c05e3a2b..5cf020df2 100644 --- a/nodes.py +++ b/nodes.py @@ -1876,6 +1876,7 @@ def load_custom_node(module_path, ignore=set()): sp = os.path.splitext(module_path) module_name = sp[0] try: + logging.debug("Trying to load custom node {}".format(module_path)) if os.path.isfile(module_path): module_spec = importlib.util.spec_from_file_location(module_name, module_path) module_dir = os.path.split(module_path)[0] From 575acb69e46a4c24ddba4e0bdb895d6e85dc9354 Mon Sep 17 00:00:00 2001 From: comfyanonymous Date: Sun, 31 Mar 2024 01:25:16 -0400 Subject: [PATCH 03/29] IP2P model loading support. This is the code to load the model and inference it with only a text prompt. This commit does not contain the nodes to properly use it with an image input. This supports both the original SD1 instructpix2pix model and the diffusers SDXL one. --- comfy/model_base.py | 34 ++++++++++++++++++++++++++++++++++ comfy/model_detection.py | 14 ++++++++++---- comfy/supported_models.py | 34 +++++++++++++++++++++++++++++++++- comfy/supported_models_base.py | 8 +++++++- 4 files changed, 84 insertions(+), 6 deletions(-) diff --git a/comfy/model_base.py b/comfy/model_base.py index 6f530d2fa..898adb665 100644 --- a/comfy/model_base.py +++ b/comfy/model_base.py @@ -473,6 +473,40 @@ class SD_X4Upscaler(BaseModel): out['y'] = comfy.conds.CONDRegular(noise_level) return out +class IP2P: + def extra_conds(self, **kwargs): + out = {} + + image = kwargs.get("concat_latent_image", None) + noise = kwargs.get("noise", None) + device = kwargs["device"] + + if image is None: + image = torch.zeros_like(noise) + + if image.shape[1:] != noise.shape[1:]: + image = utils.common_upscale(image.to(device), noise.shape[-1], noise.shape[-2], "bilinear", "center") + + image = utils.resize_to_batch_size(image, noise.shape[0]) + + out['c_concat'] = comfy.conds.CONDNoiseShape(self.process_ip2p_image_in(image)) + adm = self.encode_adm(**kwargs) + if adm is not None: + out['y'] = comfy.conds.CONDRegular(adm) + return out + +class SD15_instructpix2pix(IP2P, BaseModel): + def __init__(self, model_config, model_type=ModelType.EPS, device=None): + super().__init__(model_config, model_type, device=device) + self.process_ip2p_image_in = lambda image: image + +class SDXL_instructpix2pix(IP2P, SDXL): + def __init__(self, model_config, model_type=ModelType.EPS, device=None): + super().__init__(model_config, model_type, device=device) + # self.process_ip2p_image_in = lambda image: comfy.latent_formats.SDXL().process_in(image) + self.process_ip2p_image_in = lambda image: image + + class StableCascade_C(BaseModel): def __init__(self, model_config, model_type=ModelType.STABLE_CASCADE, device=None): super().__init__(model_config, model_type, device=device, unet_model=StageC) diff --git a/comfy/model_detection.py b/comfy/model_detection.py index bddbe2a49..795af828a 100644 --- a/comfy/model_detection.py +++ b/comfy/model_detection.py @@ -182,9 +182,9 @@ def detect_unet_config(state_dict, key_prefix): return unet_config -def model_config_from_unet_config(unet_config): +def model_config_from_unet_config(unet_config, state_dict=None): for model_config in comfy.supported_models.models: - if model_config.matches(unet_config): + if model_config.matches(unet_config, state_dict): return model_config(unet_config) logging.error("no match {}".format(unet_config)) @@ -192,7 +192,7 @@ def model_config_from_unet_config(unet_config): def model_config_from_unet(state_dict, unet_key_prefix, use_base_if_no_match=False): unet_config = detect_unet_config(state_dict, unet_key_prefix) - model_config = model_config_from_unet_config(unet_config) + model_config = model_config_from_unet_config(unet_config, state_dict) if model_config is None and use_base_if_no_match: return comfy.supported_models_base.BASE(unet_config) else: @@ -321,6 +321,12 @@ def unet_config_from_diffusers_unet(state_dict, dtype=None): 'use_linear_in_transformer': True, 'context_dim': 2048, 'num_head_channels': 64, 'transformer_depth_output': [0, 0, 0, 2, 2, 2, 10, 10, 10], 'use_temporal_attention': False, 'use_temporal_resblock': False} + SDXL_diffusers_ip2p = {'use_checkpoint': False, 'image_size': 32, 'out_channels': 4, 'use_spatial_transformer': True, 'legacy': False, + 'num_classes': 'sequential', 'adm_in_channels': 2816, 'dtype': dtype, 'in_channels': 8, 'model_channels': 320, + 'num_res_blocks': [2, 2, 2], 'transformer_depth': [0, 0, 2, 2, 10, 10], 'channel_mult': [1, 2, 4], 'transformer_depth_middle': 10, + 'use_linear_in_transformer': True, 'context_dim': 2048, 'num_head_channels': 64, 'transformer_depth_output': [0, 0, 0, 2, 2, 2, 10, 10, 10], + 'use_temporal_attention': False, 'use_temporal_resblock': False} + SSD_1B = {'use_checkpoint': False, 'image_size': 32, 'out_channels': 4, 'use_spatial_transformer': True, 'legacy': False, 'num_classes': 'sequential', 'adm_in_channels': 2816, 'dtype': dtype, 'in_channels': 4, 'model_channels': 320, 'num_res_blocks': [2, 2, 2], 'transformer_depth': [0, 0, 2, 2, 4, 4], 'transformer_depth_output': [0, 0, 0, 1, 1, 2, 10, 4, 4], @@ -351,7 +357,7 @@ def unet_config_from_diffusers_unet(state_dict, dtype=None): 'context_dim': 1024, 'num_head_channels': 64, 'transformer_depth_output': [1, 1, 1, 1, 1, 1], 'use_temporal_attention': False, 'use_temporal_resblock': False, 'disable_self_attentions': [True, False, False]} - supported_models = [SDXL, SDXL_refiner, SD21, SD15, SD21_uncliph, SD21_unclipl, SDXL_mid_cnet, SDXL_small_cnet, SDXL_diffusers_inpaint, SSD_1B, Segmind_Vega, KOALA_700M, KOALA_1B, SD09_XS] + supported_models = [SDXL, SDXL_refiner, SD21, SD15, SD21_uncliph, SD21_unclipl, SDXL_mid_cnet, SDXL_small_cnet, SDXL_diffusers_inpaint, SSD_1B, Segmind_Vega, KOALA_700M, KOALA_1B, SD09_XS, SDXL_diffusers_ip2p] for unet_config in supported_models: matches = True diff --git a/comfy/supported_models.py b/comfy/supported_models.py index 5b2eb73fd..9bfe3ea1d 100644 --- a/comfy/supported_models.py +++ b/comfy/supported_models.py @@ -334,6 +334,11 @@ class Stable_Zero123(supported_models_base.BASE): "num_head_channels": -1, } + required_keys = { + "cc_projection.weight": None, + "cc_projection.bias": None, + } + clip_vision_prefix = "cond_stage_model.model.visual." latent_format = latent_formats.SD15 @@ -439,6 +444,33 @@ class Stable_Cascade_B(Stable_Cascade_C): out = model_base.StableCascade_B(self, device=device) return out +class SD15_instructpix2pix(SD15): + unet_config = { + "context_dim": 768, + "model_channels": 320, + "use_linear_in_transformer": False, + "adm_in_channels": None, + "use_temporal_attention": False, + "in_channels": 8, + } + + def get_model(self, state_dict, prefix="", device=None): + return model_base.SD15_instructpix2pix(self, device=device) + +class SDXL_instructpix2pix(SDXL): + unet_config = { + "model_channels": 320, + "use_linear_in_transformer": True, + "transformer_depth": [0, 0, 2, 2, 10, 10], + "context_dim": 2048, + "adm_in_channels": 2816, + "use_temporal_attention": False, + "in_channels": 8, + } + + def get_model(self, state_dict, prefix="", device=None): + return model_base.SDXL_instructpix2pix(self, device=device) + +models = [Stable_Zero123, SD15_instructpix2pix, SD15, SD20, SD21UnclipL, SD21UnclipH, SDXL_instructpix2pix, SDXLRefiner, SDXL, SSD1B, KOALA_700M, KOALA_1B, Segmind_Vega, SD_X4Upscaler, Stable_Cascade_C, Stable_Cascade_B, SV3D_u, SV3D_p] -models = [Stable_Zero123, SD15, SD20, SD21UnclipL, SD21UnclipH, SDXLRefiner, SDXL, SSD1B, KOALA_700M, KOALA_1B, Segmind_Vega, SD_X4Upscaler, Stable_Cascade_C, Stable_Cascade_B, SV3D_u, SV3D_p] models += [SVD_img2vid] diff --git a/comfy/supported_models_base.py b/comfy/supported_models_base.py index 4d7e25936..6196daabf 100644 --- a/comfy/supported_models_base.py +++ b/comfy/supported_models_base.py @@ -16,6 +16,8 @@ class BASE: "num_head_channels": 64, } + required_keys = {} + clip_prefix = [] clip_vision_prefix = None noise_aug_config = None @@ -28,10 +30,14 @@ class BASE: manual_cast_dtype = None @classmethod - def matches(s, unet_config): + def matches(s, unet_config, state_dict=None): for k in s.unet_config: if k not in unet_config or s.unet_config[k] != unet_config[k]: return False + if state_dict is not None: + for k in s.required_keys: + if k not in state_dict: + return False return True def model_type(self, state_dict, prefix=""): From 130646453891bea5e86c93b689fe2daf93cae35c Mon Sep 17 00:00:00 2001 From: comfyanonymous Date: Sun, 31 Mar 2024 12:50:28 -0400 Subject: [PATCH 04/29] --force-fp16 is no longer necessary on Mac. --- README.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/README.md b/README.md index a94a212ad..ba1e844b3 100644 --- a/README.md +++ b/README.md @@ -142,7 +142,7 @@ You can install ComfyUI in Apple Mac silicon (M1 or M2) with any recent macOS ve 1. Install pytorch nightly. For instructions, read the [Accelerated PyTorch training on Mac](https://developer.apple.com/metal/pytorch/) Apple Developer guide (make sure to install the latest pytorch nightly). 1. Follow the [ComfyUI manual installation](#manual-install-windows-linux) instructions for Windows and Linux. 1. Install the ComfyUI [dependencies](#dependencies). If you have another Stable Diffusion UI [you might be able to reuse the dependencies](#i-already-have-another-ui-for-stable-diffusion-installed-do-i-really-have-to-install-all-of-these-dependencies). -1. Launch ComfyUI by running `python main.py --force-fp16`. Note that --force-fp16 will only work if you installed the latest pytorch nightly. +1. Launch ComfyUI by running `python main.py` > **Note**: Remember to add your models, VAE, LoRAs etc. to the corresponding Comfy folders, as discussed in [ComfyUI manual installation](#manual-install-windows-linux). From e6482fbbfc83cd25add0532b2e4c51d305e8a232 Mon Sep 17 00:00:00 2001 From: comfyanonymous Date: Mon, 1 Apr 2024 17:23:07 -0400 Subject: [PATCH 05/29] Refactor calc_cond_uncond_batch into calc_cond_batch. calc_cond_batch can take an arbitrary amount of cond inputs. Added a calc_cond_uncond_batch wrapper with a warning so custom nodes won't break. --- comfy/samplers.py | 67 +++++++++++++++++------------------ comfy_extras/nodes_perpneg.py | 2 +- comfy_extras/nodes_sag.py | 2 +- 3 files changed, 34 insertions(+), 37 deletions(-) diff --git a/comfy/samplers.py b/comfy/samplers.py index 3678dc818..204a98f95 100644 --- a/comfy/samplers.py +++ b/comfy/samplers.py @@ -127,30 +127,23 @@ def cond_cat(c_list): return out -def calc_cond_uncond_batch(model, cond, uncond, x_in, timestep, model_options): - out_cond = torch.zeros_like(x_in) - out_count = torch.ones_like(x_in) * 1e-37 - - out_uncond = torch.zeros_like(x_in) - out_uncond_count = torch.ones_like(x_in) * 1e-37 - - COND = 0 - UNCOND = 1 - +def calc_cond_batch(model, conds, x_in, timestep, model_options): + out_conds = [] + out_counts = [] to_run = [] - for x in cond: - p = get_area_and_mult(x, x_in, timestep) - if p is None: - continue - to_run += [(p, COND)] - if uncond is not None: - for x in uncond: - p = get_area_and_mult(x, x_in, timestep) - if p is None: - continue + for i in range(len(conds)): + out_conds.append(torch.zeros_like(x_in)) + out_counts.append(torch.ones_like(x_in) * 1e-37) - to_run += [(p, UNCOND)] + cond = conds[i] + if cond is not None: + for x in cond: + p = get_area_and_mult(x, x_in, timestep) + if p is None: + continue + + to_run += [(p, i)] while len(to_run) > 0: first = to_run[0] @@ -222,22 +215,20 @@ def calc_cond_uncond_batch(model, cond, uncond, x_in, timestep, model_options): output = model_options['model_function_wrapper'](model.apply_model, {"input": input_x, "timestep": timestep_, "c": c, "cond_or_uncond": cond_or_uncond}).chunk(batch_chunks) else: output = model.apply_model(input_x, timestep_, **c).chunk(batch_chunks) - del input_x for o in range(batch_chunks): - if cond_or_uncond[o] == COND: - out_cond[:,:,area[o][2]:area[o][0] + area[o][2],area[o][3]:area[o][1] + area[o][3]] += output[o] * mult[o] - out_count[:,:,area[o][2]:area[o][0] + area[o][2],area[o][3]:area[o][1] + area[o][3]] += mult[o] - else: - out_uncond[:,:,area[o][2]:area[o][0] + area[o][2],area[o][3]:area[o][1] + area[o][3]] += output[o] * mult[o] - out_uncond_count[:,:,area[o][2]:area[o][0] + area[o][2],area[o][3]:area[o][1] + area[o][3]] += mult[o] - del mult + cond_index = cond_or_uncond[o] + out_conds[cond_index][:,:,area[o][2]:area[o][0] + area[o][2],area[o][3]:area[o][1] + area[o][3]] += output[o] * mult[o] + out_counts[cond_index][:,:,area[o][2]:area[o][0] + area[o][2],area[o][3]:area[o][1] + area[o][3]] += mult[o] - out_cond /= out_count - del out_count - out_uncond /= out_uncond_count - del out_uncond_count - return out_cond, out_uncond + for i in range(len(out_conds)): + out_conds[i] /= out_counts[i] + + return out_conds + +def calc_cond_uncond_batch(model, cond, uncond, x_in, timestep, model_options): #TODO: remove + logging.warning("WARNING: The comfy.samplers.calc_cond_uncond_batch function is deprecated please use the calc_cond_batch one instead.") + return tuple(calc_cond_batch(model, [cond, uncond], x_in, timestep, model_options)) #The main sampling function shared by all the samplers #Returns denoised @@ -247,7 +238,13 @@ def sampling_function(model, x, timestep, uncond, cond, cond_scale, model_option else: uncond_ = uncond - cond_pred, uncond_pred = calc_cond_uncond_batch(model, cond, uncond_, x, timestep, model_options) + + conds = [cond, uncond_] + + out = calc_cond_batch(model, conds, x, timestep, model_options) + cond_pred = out[0] + uncond_pred = out[1] + if "sampler_cfg_function" in model_options: args = {"cond": x - cond_pred, "uncond": x - uncond_pred, "cond_scale": cond_scale, "timestep": timestep, "input": x, "sigma": timestep, "cond_denoised": cond_pred, "uncond_denoised": uncond_pred, "model": model, "model_options": model_options} diff --git a/comfy_extras/nodes_perpneg.py b/comfy_extras/nodes_perpneg.py index dc73c5528..9e8a218f9 100644 --- a/comfy_extras/nodes_perpneg.py +++ b/comfy_extras/nodes_perpneg.py @@ -31,7 +31,7 @@ class PerpNeg: model_options = args["model_options"] nocond_processed = comfy.samplers.encode_model_conds(model.extra_conds, nocond, x, x.device, "negative") - (noise_pred_nocond, _) = comfy.samplers.calc_cond_uncond_batch(model, nocond_processed, None, x, sigma, model_options) + (noise_pred_nocond,) = comfy.samplers.calc_cond_batch(model, [nocond_processed], x, sigma, model_options) pos = noise_pred_pos - noise_pred_nocond neg = noise_pred_neg - noise_pred_nocond diff --git a/comfy_extras/nodes_sag.py b/comfy_extras/nodes_sag.py index bbd380807..69084e91d 100644 --- a/comfy_extras/nodes_sag.py +++ b/comfy_extras/nodes_sag.py @@ -150,7 +150,7 @@ class SelfAttentionGuidance: degraded = create_blur_map(uncond_pred, uncond_attn, sag_sigma, sag_threshold) degraded_noised = degraded + x - uncond_pred # call into the UNet - (sag, _) = comfy.samplers.calc_cond_uncond_batch(model, uncond, None, degraded_noised, sigma, model_options) + (sag,) = comfy.samplers.calc_cond_batch(model, [uncond], degraded_noised, sigma, model_options) return cfg_result + (degraded - sag) * sag_scale m.set_model_sampler_post_cfg_function(post_cfg_function, disable_cfg1_optimization=True) From 6c6a39251fe313c56a88c90d820009073b623dfe Mon Sep 17 00:00:00 2001 From: comfyanonymous Date: Tue, 2 Apr 2024 11:46:34 -0400 Subject: [PATCH 06/29] Fix saving text encoder in fp8. --- comfy/diffusers_convert.py | 19 +++++++++++++++++-- 1 file changed, 17 insertions(+), 2 deletions(-) diff --git a/comfy/diffusers_convert.py b/comfy/diffusers_convert.py index 08018c54d..ed2a45fea 100644 --- a/comfy/diffusers_convert.py +++ b/comfy/diffusers_convert.py @@ -206,6 +206,21 @@ textenc_pattern = re.compile("|".join(protected.keys())) # Ordering is from https://github.com/pytorch/pytorch/blob/master/test/cpp/api/modules.cpp code2idx = {"q": 0, "k": 1, "v": 2} +# This function exists because at the time of writing torch.cat can't do fp8 with cuda +def cat_tensors(tensors): + x = 0 + for t in tensors: + x += t.shape[0] + + shape = [x] + list(tensors[0].shape)[1:] + out = torch.empty(shape, device=tensors[0].device, dtype=tensors[0].dtype) + + x = 0 + for t in tensors: + out[x:x + t.shape[0]] = t + x += t.shape[0] + + return out def convert_text_enc_state_dict_v20(text_enc_dict, prefix=""): new_state_dict = {} @@ -249,13 +264,13 @@ def convert_text_enc_state_dict_v20(text_enc_dict, prefix=""): if None in tensors: raise Exception("CORRUPTED MODEL: one of the q-k-v values for the text encoder was missing") relabelled_key = textenc_pattern.sub(lambda m: protected[re.escape(m.group(0))], k_pre) - new_state_dict[relabelled_key + ".in_proj_weight"] = torch.cat(tensors) + new_state_dict[relabelled_key + ".in_proj_weight"] = cat_tensors(tensors) for k_pre, tensors in capture_qkv_bias.items(): if None in tensors: raise Exception("CORRUPTED MODEL: one of the q-k-v values for the text encoder was missing") relabelled_key = textenc_pattern.sub(lambda m: protected[re.escape(m.group(0))], k_pre) - new_state_dict[relabelled_key + ".in_proj_bias"] = torch.cat(tensors) + new_state_dict[relabelled_key + ".in_proj_bias"] = cat_tensors(tensors) return new_state_dict From 57753c964affd18d2b87d2a47fe6b375bca39004 Mon Sep 17 00:00:00 2001 From: comfyanonymous Date: Wed, 3 Apr 2024 16:34:19 -0400 Subject: [PATCH 07/29] Refactor sampling code for more advanced sampler nodes. --- comfy/sample.py | 38 +++++++++++++++--------- comfy/samplers.py | 74 +++++++++++++++++++++++++++-------------------- 2 files changed, 68 insertions(+), 44 deletions(-) diff --git a/comfy/sample.py b/comfy/sample.py index 5c8a7d130..3c65d0a85 100644 --- a/comfy/sample.py +++ b/comfy/sample.py @@ -52,9 +52,16 @@ def convert_cond(cond): out.append(temp) return out -def get_additional_models(positive, negative, dtype): - """loads additional models in positive and negative conditioning""" - control_nets = set(get_models_from_cond(positive, "control") + get_models_from_cond(negative, "control")) +def get_additional_models(conds, dtype): + """loads additional models in conditioning""" + cnets = [] + gligen = [] + + for i in range(len(conds)): + cnets += get_models_from_cond(conds[i], "control") + gligen += get_models_from_cond(conds[i], "gligen") + + control_nets = set(cnets) inference_memory = 0 control_models = [] @@ -62,7 +69,6 @@ def get_additional_models(positive, negative, dtype): control_models += m.get_models() inference_memory += m.inference_memory_requirements(dtype) - gligen = get_models_from_cond(positive, "gligen") + get_models_from_cond(negative, "gligen") gligen = [x[1] for x in gligen] models = control_models + gligen return models, inference_memory @@ -73,24 +79,25 @@ def cleanup_additional_models(models): if hasattr(m, 'cleanup'): m.cleanup() -def prepare_sampling(model, noise_shape, positive, negative, noise_mask): +def prepare_sampling(model, noise_shape, conds, noise_mask): device = model.load_device - positive = convert_cond(positive) - negative = convert_cond(negative) + for i in range(len(conds)): + conds[i] = convert_cond(conds[i]) if noise_mask is not None: noise_mask = prepare_mask(noise_mask, noise_shape, device) real_model = None - models, inference_memory = get_additional_models(positive, negative, model.model_dtype()) + models, inference_memory = get_additional_models(conds, model.model_dtype()) comfy.model_management.load_models_gpu([model] + models, model.memory_required([noise_shape[0] * 2] + list(noise_shape[1:])) + inference_memory) real_model = model.model - return real_model, positive, negative, noise_mask, models + return real_model, conds, noise_mask, models def sample(model, noise, steps, cfg, sampler_name, scheduler, positive, negative, latent_image, denoise=1.0, disable_noise=False, start_step=None, last_step=None, force_full_denoise=False, noise_mask=None, sigmas=None, callback=None, disable_pbar=False, seed=None): - real_model, positive_copy, negative_copy, noise_mask, models = prepare_sampling(model, noise.shape, positive, negative, noise_mask) + real_model, conds_copy, noise_mask, models = prepare_sampling(model, noise.shape, [positive, negative], noise_mask) + positive_copy, negative_copy = conds_copy noise = noise.to(model.load_device) latent_image = latent_image.to(model.load_device) @@ -105,14 +112,19 @@ def sample(model, noise, steps, cfg, sampler_name, scheduler, positive, negative return samples def sample_custom(model, noise, cfg, sampler, sigmas, positive, negative, latent_image, noise_mask=None, callback=None, disable_pbar=False, seed=None): - real_model, positive_copy, negative_copy, noise_mask, models = prepare_sampling(model, noise.shape, positive, negative, noise_mask) + real_model, conds, noise_mask, models = prepare_sampling(model, noise.shape, [positive, negative], noise_mask) noise = noise.to(model.load_device) latent_image = latent_image.to(model.load_device) sigmas = sigmas.to(model.load_device) - samples = comfy.samplers.sample(real_model, noise, positive_copy, negative_copy, cfg, model.load_device, sampler, sigmas, model_options=model.model_options, latent_image=latent_image, denoise_mask=noise_mask, callback=callback, disable_pbar=disable_pbar, seed=seed) + samples = comfy.samplers.sample(real_model, noise, conds[0], conds[1], cfg, model.load_device, sampler, sigmas, model_options=model.model_options, latent_image=latent_image, denoise_mask=noise_mask, callback=callback, disable_pbar=disable_pbar, seed=seed) samples = samples.to(comfy.model_management.intermediate_device()) cleanup_additional_models(models) - cleanup_additional_models(set(get_models_from_cond(positive_copy, "control") + get_models_from_cond(negative_copy, "control"))) + + control_cleanup = [] + for i in range(len(conds)): + control_cleanup += get_models_from_cond(conds[i], "control") + + cleanup_additional_models(set(control_cleanup)) return samples diff --git a/comfy/samplers.py b/comfy/samplers.py index 204a98f95..f18de200e 100644 --- a/comfy/samplers.py +++ b/comfy/samplers.py @@ -260,11 +260,12 @@ def sampling_function(model, x, timestep, uncond, cond, cond_scale, model_option return cfg_result class CFGNoisePredictor(torch.nn.Module): - def __init__(self, model): + def __init__(self, model, cond_scale=1.0): super().__init__() self.inner_model = model - def apply_model(self, x, timestep, cond, uncond, cond_scale, model_options={}, seed=None): - out = sampling_function(self.inner_model, x, timestep, uncond, cond, cond_scale, model_options=model_options, seed=seed) + self.cond_scale = cond_scale + def apply_model(self, x, timestep, conds, model_options={}, seed=None): + out = sampling_function(self.inner_model, x, timestep, conds.get("negative", None), conds.get("positive", None), self.cond_scale, model_options=model_options, seed=seed) return out def forward(self, *args, **kwargs): return self.apply_model(*args, **kwargs) @@ -274,13 +275,13 @@ class KSamplerX0Inpaint(torch.nn.Module): super().__init__() self.inner_model = model self.sigmas = sigmas - def forward(self, x, sigma, uncond, cond, cond_scale, denoise_mask, model_options={}, seed=None): + def forward(self, x, sigma, conds, denoise_mask, model_options={}, seed=None): if denoise_mask is not None: if "denoise_mask_function" in model_options: denoise_mask = model_options["denoise_mask_function"](sigma, denoise_mask, extra_options={"model": self.inner_model, "sigmas": self.sigmas}) latent_mask = 1. - denoise_mask x = x * denoise_mask + self.inner_model.inner_model.model_sampling.noise_scaling(sigma.reshape([sigma.shape[0]] + [1] * (len(self.noise.shape) - 1)), self.noise, self.latent_image) * latent_mask - out = self.inner_model(x, sigma, cond=cond, uncond=uncond, cond_scale=cond_scale, model_options=model_options, seed=seed) + out = self.inner_model(x, sigma, conds=conds, model_options=model_options, seed=seed) if denoise_mask is not None: out = out * denoise_mask + self.latent_image * latent_mask return out @@ -568,45 +569,56 @@ def ksampler(sampler_name, extra_options={}, inpaint_options={}): return KSAMPLER(sampler_function, extra_options, inpaint_options) -def wrap_model(model): - model_denoise = CFGNoisePredictor(model) - return model_denoise -def sample(model, noise, positive, negative, cfg, device, sampler, sigmas, model_options={}, latent_image=None, denoise_mask=None, callback=None, disable_pbar=False, seed=None): - positive = positive[:] - negative = negative[:] +def process_conds(model, noise, conds, device, latent_image=None, denoise_mask=None, seed=None): + for k in conds: + conds[k] = conds[k][:] + resolve_areas_and_cond_masks(conds[k], noise.shape[2], noise.shape[3], device) - resolve_areas_and_cond_masks(positive, noise.shape[2], noise.shape[3], device) - resolve_areas_and_cond_masks(negative, noise.shape[2], noise.shape[3], device) + for k in conds: + calculate_start_end_timesteps(model, conds[k]) - model_wrap = wrap_model(model) + if hasattr(model, 'extra_conds'): + for k in conds: + conds[k] = encode_model_conds(model.extra_conds, conds[k], noise, device, k, latent_image=latent_image, denoise_mask=denoise_mask, seed=seed) - calculate_start_end_timesteps(model, negative) - calculate_start_end_timesteps(model, positive) + #make sure each cond area has an opposite one with the same area + for k in conds: + for c in conds[k]: + for kk in conds: + if k != kk: + create_cond_with_same_area_if_none(conds[kk], c) + for k in conds: + pre_run_control(model, conds[k]) + + if "positive" in conds: + positive = conds["positive"] + for k in conds: + if k != "positive": + apply_empty_x_to_equal_area(list(filter(lambda c: c.get('control_apply_to_uncond', False) == True, positive)), conds[k], 'control', lambda cond_cnets, x: cond_cnets[x]) + apply_empty_x_to_equal_area(positive, conds[k], 'gligen', lambda cond_cnets, x: cond_cnets[x]) + + return conds + + +def sample_advanced(model, noise, conds, guider_class, device, sampler, sigmas, model_options={}, latent_image=None, denoise_mask=None, callback=None, disable_pbar=False, seed=None): if latent_image is not None and torch.count_nonzero(latent_image) > 0: #Don't shift the empty latent image. latent_image = model.process_latent_in(latent_image) - if hasattr(model, 'extra_conds'): - positive = encode_model_conds(model.extra_conds, positive, noise, device, "positive", latent_image=latent_image, denoise_mask=denoise_mask, seed=seed) - negative = encode_model_conds(model.extra_conds, negative, noise, device, "negative", latent_image=latent_image, denoise_mask=denoise_mask, seed=seed) + conds = process_conds(model, noise, conds, device, latent_image, denoise_mask, seed) + model_wrap = guider_class(model) - #make sure each cond area has an opposite one with the same area - for c in positive: - create_cond_with_same_area_if_none(negative, c) - for c in negative: - create_cond_with_same_area_if_none(positive, c) - - pre_run_control(model, negative + positive) - - apply_empty_x_to_equal_area(list(filter(lambda c: c.get('control_apply_to_uncond', False) == True, positive)), negative, 'control', lambda cond_cnets, x: cond_cnets[x]) - apply_empty_x_to_equal_area(positive, negative, 'gligen', lambda cond_cnets, x: cond_cnets[x]) - - extra_args = {"cond":positive, "uncond":negative, "cond_scale": cfg, "model_options": model_options, "seed":seed} + extra_args = {"conds": conds, "model_options": model_options, "seed":seed} samples = sampler.sample(model_wrap, sigmas, extra_args, callback, noise, latent_image, denoise_mask, disable_pbar) return model.process_latent_out(samples.to(torch.float32)) + +def sample(model, noise, positive, negative, cfg, device, sampler, sigmas, model_options={}, latent_image=None, denoise_mask=None, callback=None, disable_pbar=False, seed=None): + return sample_advanced(model, noise, {"positive": positive, "negative": negative}, lambda a: CFGNoisePredictor(a, cfg), device, sampler, sigmas, model_options, latent_image, denoise_mask, callback, disable_pbar, seed) + + SCHEDULER_NAMES = ["normal", "karras", "exponential", "sgm_uniform", "simple", "ddim_uniform"] SAMPLER_NAMES = KSAMPLER_NAMES + ["ddim", "uni_pc", "uni_pc_bh2"] From 0542088ef895b4825df80fd3babf91513441af65 Mon Sep 17 00:00:00 2001 From: comfyanonymous Date: Thu, 4 Apr 2024 00:48:42 -0400 Subject: [PATCH 08/29] Refactor sampler code for more advanced sampler nodes part 2. --- comfy/sample.py | 102 +++--------------------------- comfy/sampler_helpers.py | 76 +++++++++++++++++++++++ comfy/samplers.py | 130 ++++++++++++++++++++++++--------------- 3 files changed, 165 insertions(+), 143 deletions(-) create mode 100644 comfy/sampler_helpers.py diff --git a/comfy/sample.py b/comfy/sample.py index 3c65d0a85..e51bd67d6 100644 --- a/comfy/sample.py +++ b/comfy/sample.py @@ -1,10 +1,9 @@ import torch import comfy.model_management import comfy.samplers -import comfy.conds import comfy.utils -import math import numpy as np +import logging def prepare_noise(latent_image, seed, noise_inds=None): """ @@ -25,106 +24,21 @@ def prepare_noise(latent_image, seed, noise_inds=None): noises = torch.cat(noises, axis=0) return noises -def prepare_mask(noise_mask, shape, device): - """ensures noise mask is of proper dimensions""" - noise_mask = torch.nn.functional.interpolate(noise_mask.reshape((-1, 1, noise_mask.shape[-2], noise_mask.shape[-1])), size=(shape[2], shape[3]), mode="bilinear") - noise_mask = torch.cat([noise_mask] * shape[1], dim=1) - noise_mask = comfy.utils.repeat_to_batch_size(noise_mask, shape[0]) - noise_mask = noise_mask.to(device) - return noise_mask - -def get_models_from_cond(cond, model_type): - models = [] - for c in cond: - if model_type in c: - models += [c[model_type]] - return models - -def convert_cond(cond): - out = [] - for c in cond: - temp = c[1].copy() - model_conds = temp.get("model_conds", {}) - if c[0] is not None: - model_conds["c_crossattn"] = comfy.conds.CONDCrossAttn(c[0]) #TODO: remove - temp["cross_attn"] = c[0] - temp["model_conds"] = model_conds - out.append(temp) - return out - -def get_additional_models(conds, dtype): - """loads additional models in conditioning""" - cnets = [] - gligen = [] - - for i in range(len(conds)): - cnets += get_models_from_cond(conds[i], "control") - gligen += get_models_from_cond(conds[i], "gligen") - - control_nets = set(cnets) - - inference_memory = 0 - control_models = [] - for m in control_nets: - control_models += m.get_models() - inference_memory += m.inference_memory_requirements(dtype) - - gligen = [x[1] for x in gligen] - models = control_models + gligen - return models, inference_memory +def prepare_sampling(model, noise_shape, positive, negative, noise_mask): + logging.warning("Warning: comfy.sample.prepare_sampling isn't used anymore and can be removed") + return model, positive, negative, noise_mask, [] def cleanup_additional_models(models): - """cleanup additional models that were loaded""" - for m in models: - if hasattr(m, 'cleanup'): - m.cleanup() - -def prepare_sampling(model, noise_shape, conds, noise_mask): - device = model.load_device - for i in range(len(conds)): - conds[i] = convert_cond(conds[i]) - - if noise_mask is not None: - noise_mask = prepare_mask(noise_mask, noise_shape, device) - - real_model = None - models, inference_memory = get_additional_models(conds, model.model_dtype()) - comfy.model_management.load_models_gpu([model] + models, model.memory_required([noise_shape[0] * 2] + list(noise_shape[1:])) + inference_memory) - real_model = model.model - - return real_model, conds, noise_mask, models - + logging.warning("Warning: comfy.sample.cleanup_additional_models isn't used anymore and can be removed") def sample(model, noise, steps, cfg, sampler_name, scheduler, positive, negative, latent_image, denoise=1.0, disable_noise=False, start_step=None, last_step=None, force_full_denoise=False, noise_mask=None, sigmas=None, callback=None, disable_pbar=False, seed=None): - real_model, conds_copy, noise_mask, models = prepare_sampling(model, noise.shape, [positive, negative], noise_mask) - positive_copy, negative_copy = conds_copy + sampler = comfy.samplers.KSampler(model, steps=steps, device=model.load_device, sampler=sampler_name, scheduler=scheduler, denoise=denoise, model_options=model.model_options) - noise = noise.to(model.load_device) - latent_image = latent_image.to(model.load_device) - - sampler = comfy.samplers.KSampler(real_model, steps=steps, device=model.load_device, sampler=sampler_name, scheduler=scheduler, denoise=denoise, model_options=model.model_options) - - samples = sampler.sample(noise, positive_copy, negative_copy, cfg=cfg, latent_image=latent_image, start_step=start_step, last_step=last_step, force_full_denoise=force_full_denoise, denoise_mask=noise_mask, sigmas=sigmas, callback=callback, disable_pbar=disable_pbar, seed=seed) + samples = sampler.sample(noise, positive, negative, cfg=cfg, latent_image=latent_image, start_step=start_step, last_step=last_step, force_full_denoise=force_full_denoise, denoise_mask=noise_mask, sigmas=sigmas, callback=callback, disable_pbar=disable_pbar, seed=seed) samples = samples.to(comfy.model_management.intermediate_device()) - - cleanup_additional_models(models) - cleanup_additional_models(set(get_models_from_cond(positive_copy, "control") + get_models_from_cond(negative_copy, "control"))) return samples def sample_custom(model, noise, cfg, sampler, sigmas, positive, negative, latent_image, noise_mask=None, callback=None, disable_pbar=False, seed=None): - real_model, conds, noise_mask, models = prepare_sampling(model, noise.shape, [positive, negative], noise_mask) - noise = noise.to(model.load_device) - latent_image = latent_image.to(model.load_device) - sigmas = sigmas.to(model.load_device) - - samples = comfy.samplers.sample(real_model, noise, conds[0], conds[1], cfg, model.load_device, sampler, sigmas, model_options=model.model_options, latent_image=latent_image, denoise_mask=noise_mask, callback=callback, disable_pbar=disable_pbar, seed=seed) + samples = comfy.samplers.sample(model, noise, positive, negative, cfg, model.load_device, sampler, sigmas, model_options=model.model_options, latent_image=latent_image, denoise_mask=noise_mask, callback=callback, disable_pbar=disable_pbar, seed=seed) samples = samples.to(comfy.model_management.intermediate_device()) - cleanup_additional_models(models) - - control_cleanup = [] - for i in range(len(conds)): - control_cleanup += get_models_from_cond(conds[i], "control") - - cleanup_additional_models(set(control_cleanup)) return samples - diff --git a/comfy/sampler_helpers.py b/comfy/sampler_helpers.py new file mode 100644 index 000000000..a18abd9e9 --- /dev/null +++ b/comfy/sampler_helpers.py @@ -0,0 +1,76 @@ +import torch +import comfy.model_management +import comfy.conds + +def prepare_mask(noise_mask, shape, device): + """ensures noise mask is of proper dimensions""" + noise_mask = torch.nn.functional.interpolate(noise_mask.reshape((-1, 1, noise_mask.shape[-2], noise_mask.shape[-1])), size=(shape[2], shape[3]), mode="bilinear") + noise_mask = torch.cat([noise_mask] * shape[1], dim=1) + noise_mask = comfy.utils.repeat_to_batch_size(noise_mask, shape[0]) + noise_mask = noise_mask.to(device) + return noise_mask + +def get_models_from_cond(cond, model_type): + models = [] + for c in cond: + if model_type in c: + models += [c[model_type]] + return models + +def convert_cond(cond): + out = [] + for c in cond: + temp = c[1].copy() + model_conds = temp.get("model_conds", {}) + if c[0] is not None: + model_conds["c_crossattn"] = comfy.conds.CONDCrossAttn(c[0]) #TODO: remove + temp["cross_attn"] = c[0] + temp["model_conds"] = model_conds + out.append(temp) + return out + +def get_additional_models(conds, dtype): + """loads additional models in conditioning""" + cnets = [] + gligen = [] + + for k in conds: + cnets += get_models_from_cond(conds[k], "control") + gligen += get_models_from_cond(conds[k], "gligen") + + control_nets = set(cnets) + + inference_memory = 0 + control_models = [] + for m in control_nets: + control_models += m.get_models() + inference_memory += m.inference_memory_requirements(dtype) + + gligen = [x[1] for x in gligen] + models = control_models + gligen + return models, inference_memory + +def cleanup_additional_models(models): + """cleanup additional models that were loaded""" + for m in models: + if hasattr(m, 'cleanup'): + m.cleanup() + + +def prepare_sampling(model, noise_shape, conds): + device = model.load_device + real_model = None + models, inference_memory = get_additional_models(conds, model.model_dtype()) + comfy.model_management.load_models_gpu([model] + models, model.memory_required([noise_shape[0] * 2] + list(noise_shape[1:])) + inference_memory) + real_model = model.model + + return real_model, conds, models + +def cleanup_models(conds, models): + cleanup_additional_models(models) + + control_cleanup = [] + for k in conds: + control_cleanup += get_models_from_cond(conds[k], "control") + + cleanup_additional_models(set(control_cleanup)) diff --git a/comfy/samplers.py b/comfy/samplers.py index f18de200e..57f5632c5 100644 --- a/comfy/samplers.py +++ b/comfy/samplers.py @@ -5,6 +5,7 @@ import collections from comfy import model_management import math import logging +import comfy.sampler_helpers def get_area_and_mult(conds, x_in, timestep_in): area = (x_in.shape[2], x_in.shape[3], 0, 0) @@ -230,58 +231,45 @@ def calc_cond_uncond_batch(model, cond, uncond, x_in, timestep, model_options): logging.warning("WARNING: The comfy.samplers.calc_cond_uncond_batch function is deprecated please use the calc_cond_batch one instead.") return tuple(calc_cond_batch(model, [cond, uncond], x_in, timestep, model_options)) +def cfg_function(model, cond_pred, uncond_pred, cond_scale, x, timestep, model_options={}): + if "sampler_cfg_function" in model_options: + args = {"cond": x - cond_pred, "uncond": x - uncond_pred, "cond_scale": cond_scale, "timestep": timestep, "input": x, "sigma": timestep, + "cond_denoised": cond_pred, "uncond_denoised": uncond_pred, "model": model, "model_options": model_options} + cfg_result = x - model_options["sampler_cfg_function"](args) + else: + cfg_result = uncond_pred + (cond_pred - uncond_pred) * cond_scale + + for fn in model_options.get("sampler_post_cfg_function", []): + args = {"denoised": cfg_result, "cond": cond, "uncond": uncond, "model": model, "uncond_denoised": uncond_pred, "cond_denoised": cond_pred, + "sigma": timestep, "model_options": model_options, "input": x} + cfg_result = fn(args) + + return cfg_result + #The main sampling function shared by all the samplers #Returns denoised def sampling_function(model, x, timestep, uncond, cond, cond_scale, model_options={}, seed=None): - if math.isclose(cond_scale, 1.0) and model_options.get("disable_cfg1_optimization", False) == False: - uncond_ = None - else: - uncond_ = uncond + if math.isclose(cond_scale, 1.0) and model_options.get("disable_cfg1_optimization", False) == False: + uncond_ = None + else: + uncond_ = uncond + + conds = [cond, uncond_] + out = calc_cond_batch(model, conds, x, timestep, model_options) + return cfg_function(model, out[0], out[1], cond_scale, x, timestep, model_options=model_options) - conds = [cond, uncond_] - - out = calc_cond_batch(model, conds, x, timestep, model_options) - cond_pred = out[0] - uncond_pred = out[1] - - if "sampler_cfg_function" in model_options: - args = {"cond": x - cond_pred, "uncond": x - uncond_pred, "cond_scale": cond_scale, "timestep": timestep, "input": x, "sigma": timestep, - "cond_denoised": cond_pred, "uncond_denoised": uncond_pred, "model": model, "model_options": model_options} - cfg_result = x - model_options["sampler_cfg_function"](args) - else: - cfg_result = uncond_pred + (cond_pred - uncond_pred) * cond_scale - - for fn in model_options.get("sampler_post_cfg_function", []): - args = {"denoised": cfg_result, "cond": cond, "uncond": uncond, "model": model, "uncond_denoised": uncond_pred, "cond_denoised": cond_pred, - "sigma": timestep, "model_options": model_options, "input": x} - cfg_result = fn(args) - - return cfg_result - -class CFGNoisePredictor(torch.nn.Module): - def __init__(self, model, cond_scale=1.0): - super().__init__() - self.inner_model = model - self.cond_scale = cond_scale - def apply_model(self, x, timestep, conds, model_options={}, seed=None): - out = sampling_function(self.inner_model, x, timestep, conds.get("negative", None), conds.get("positive", None), self.cond_scale, model_options=model_options, seed=seed) - return out - def forward(self, *args, **kwargs): - return self.apply_model(*args, **kwargs) - -class KSamplerX0Inpaint(torch.nn.Module): +class KSamplerX0Inpaint: def __init__(self, model, sigmas): - super().__init__() self.inner_model = model self.sigmas = sigmas - def forward(self, x, sigma, conds, denoise_mask, model_options={}, seed=None): + def __call__(self, x, sigma, denoise_mask, model_options={}, seed=None): if denoise_mask is not None: if "denoise_mask_function" in model_options: denoise_mask = model_options["denoise_mask_function"](sigma, denoise_mask, extra_options={"model": self.inner_model, "sigmas": self.sigmas}) latent_mask = 1. - denoise_mask x = x * denoise_mask + self.inner_model.inner_model.model_sampling.noise_scaling(sigma.reshape([sigma.shape[0]] + [1] * (len(self.noise.shape) - 1)), self.noise, self.latent_image) * latent_mask - out = self.inner_model(x, sigma, conds=conds, model_options=model_options, seed=seed) + out = self.inner_model(x, sigma, model_options=model_options, seed=seed) if denoise_mask is not None: out = out * denoise_mask + self.latent_image * latent_mask return out @@ -601,22 +589,66 @@ def process_conds(model, noise, conds, device, latent_image=None, denoise_mask=N return conds +class CFGGuider: + def __init__(self, model_patcher): + self.model_patcher = model_patcher + self.model_options = model_patcher.model_options + self.original_conds = {} + self.cfg = 1.0 -def sample_advanced(model, noise, conds, guider_class, device, sampler, sigmas, model_options={}, latent_image=None, denoise_mask=None, callback=None, disable_pbar=False, seed=None): - if latent_image is not None and torch.count_nonzero(latent_image) > 0: #Don't shift the empty latent image. - latent_image = model.process_latent_in(latent_image) + def set_conds(self, conds): + for k in conds: + self.original_conds[k] = comfy.sampler_helpers.convert_cond(conds[k]) - conds = process_conds(model, noise, conds, device, latent_image, denoise_mask, seed) - model_wrap = guider_class(model) + def set_cfg(self, cfg): + self.cfg = cfg - extra_args = {"conds": conds, "model_options": model_options, "seed":seed} + def __call__(self, *args, **kwargs): + return self.predict_noise(*args, **kwargs) - samples = sampler.sample(model_wrap, sigmas, extra_args, callback, noise, latent_image, denoise_mask, disable_pbar) - return model.process_latent_out(samples.to(torch.float32)) + def predict_noise(self, x, timestep, model_options={}, seed=None): + return sampling_function(self.inner_model, x, timestep, self.conds.get("negative", None), self.conds.get("positive", None), self.cfg, model_options=model_options, seed=seed) + + def inner_sample(self, noise, latent_image, device, sampler, sigmas, denoise_mask, callback, disable_pbar, seed): + if latent_image is not None and torch.count_nonzero(latent_image) > 0: #Don't shift the empty latent image. + latent_image = self.inner_model.process_latent_in(latent_image) + + self.conds = process_conds(self.inner_model, noise, self.conds, device, latent_image, denoise_mask, seed) + + extra_args = {"model_options": self.model_options, "seed":seed} + + samples = sampler.sample(self, sigmas, extra_args, callback, noise, latent_image, denoise_mask, disable_pbar) + return self.inner_model.process_latent_out(samples.to(torch.float32)) + + def sample(self, noise, latent_image, sampler, sigmas, denoise_mask=None, callback=None, disable_pbar=False, seed=None): + self.conds = {} + for k in self.original_conds: + self.conds[k] = list(map(lambda a: a.copy(), self.original_conds[k])) + + self.inner_model, self.conds, self.loaded_models = comfy.sampler_helpers.prepare_sampling(self.model_patcher, noise.shape, self.conds) + device = self.model_patcher.load_device + + if denoise_mask is not None: + denoise_mask = comfy.sampler_helpers.prepare_mask(denoise_mask, noise.shape, device) + + noise = noise.to(device) + latent_image = latent_image.to(device) + sigmas = sigmas.to(device) + + output = self.inner_sample(noise, latent_image, device, sampler, sigmas, denoise_mask, callback, disable_pbar, seed) + + comfy.sampler_helpers.cleanup_models(self.conds, self.loaded_models) + del self.inner_model + del self.conds + del self.loaded_models + return output def sample(model, noise, positive, negative, cfg, device, sampler, sigmas, model_options={}, latent_image=None, denoise_mask=None, callback=None, disable_pbar=False, seed=None): - return sample_advanced(model, noise, {"positive": positive, "negative": negative}, lambda a: CFGNoisePredictor(a, cfg), device, sampler, sigmas, model_options, latent_image, denoise_mask, callback, disable_pbar, seed) + cfg_guider = CFGGuider(model) + cfg_guider.set_conds({"positive": positive, "negative": negative}) + cfg_guider.set_cfg(cfg) + return cfg_guider.sample(noise, latent_image, sampler, sigmas, denoise_mask, callback, disable_pbar, seed) SCHEDULER_NAMES = ["normal", "karras", "exponential", "sgm_uniform", "simple", "ddim_uniform"] @@ -676,7 +708,7 @@ class KSampler: steps += 1 discard_penultimate_sigma = True - sigmas = calculate_sigmas_scheduler(self.model, self.scheduler, steps) + sigmas = calculate_sigmas_scheduler(self.model.model, self.scheduler, steps) if discard_penultimate_sigma: sigmas = torch.cat([sigmas[:-2], sigmas[-1:]]) From f117566299edd621ea8b00b70cf81d5d96c20917 Mon Sep 17 00:00:00 2001 From: comfyanonymous Date: Thu, 4 Apr 2024 01:32:25 -0400 Subject: [PATCH 09/29] SamplerCustomAdvanced node. This node enables the creation of nodes to change the guider/denoiser and the noise algorithm. --- comfy_extras/nodes_custom_sampler.py | 122 ++++++++++++++++++++++++++- 1 file changed, 119 insertions(+), 3 deletions(-) diff --git a/comfy_extras/nodes_custom_sampler.py b/comfy_extras/nodes_custom_sampler.py index 72ff7957f..32b2a56d6 100644 --- a/comfy_extras/nodes_custom_sampler.py +++ b/comfy_extras/nodes_custom_sampler.py @@ -310,6 +310,24 @@ class SamplerDPMAdaptative: "s_noise":s_noise }) return (sampler, ) +class Noise_EmptyNoise: + def __init__(self): + self.seed = 0 + + def generate_noise(self, input_latent): + latent_image = input_latent["samples"] + return torch.zeros(shape, dtype=latent_image.dtype, layout=latent_image.layout, device="cpu") + + +class Noise_RandomNoise: + def __init__(self, seed): + self.seed = seed + + def generate_noise(self, input_latent): + latent_image = input_latent["samples"] + batch_inds = input_latent["batch_index"] if "batch_index" in input_latent else None + return comfy.sample.prepare_noise(latent_image, self.seed, batch_inds) + class SamplerCustom: @classmethod def INPUT_TYPES(s): @@ -337,10 +355,9 @@ class SamplerCustom: latent = latent_image latent_image = latent["samples"] if not add_noise: - noise = torch.zeros(latent_image.size(), dtype=latent_image.dtype, layout=latent_image.layout, device="cpu") + noise = Noise_EmptyNoise().generate_noise(latent) else: - batch_inds = latent["batch_index"] if "batch_index" in latent else None - noise = comfy.sample.prepare_noise(latent_image, noise_seed, batch_inds) + noise = Noise_RandomNoise(noise_seed).generate_noise(latent) noise_mask = None if "noise_mask" in latent: @@ -361,6 +378,100 @@ class SamplerCustom: out_denoised = out return (out, out_denoised) + +class CFGGuider: + @classmethod + def INPUT_TYPES(s): + return {"required": + {"model": ("MODEL",), + "positive": ("CONDITIONING", ), + "negative": ("CONDITIONING", ), + "cfg": ("FLOAT", {"default": 8.0, "min": 0.0, "max": 100.0, "step":0.1, "round": 0.01}), + } + } + + RETURN_TYPES = ("GUIDER",) + + FUNCTION = "get_guider" + CATEGORY = "sampling/custom_sampling/guiders" + + def get_guider(self, model, positive, negative, cfg): + guider = comfy.samplers.CFGGuider(model) + guider.set_conds({"positive": positive, "negative": negative}) + guider.set_cfg(cfg) + return (guider,) + + +class DisableNoise: + @classmethod + def INPUT_TYPES(s): + return {"required":{ + } + } + + RETURN_TYPES = ("NOISE",) + FUNCTION = "get_noise" + CATEGORY = "sampling/custom_sampling/noise" + + def get_noise(self, noise_seed): + return (Noise_EmptyNoise(),) + + +class RandomNoise(DisableNoise): + @classmethod + def INPUT_TYPES(s): + return {"required":{ + "noise_seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff}), + } + } + + def get_noise(self, noise_seed): + return (Noise_RandomNoise(noise_seed),) + + +class SamplerCustomAdvanced: + @classmethod + def INPUT_TYPES(s): + return {"required": + {"noise": ("NOISE", ), + "guider": ("GUIDER", ), + "sampler": ("SAMPLER", ), + "sigmas": ("SIGMAS", ), + "latent_image": ("LATENT", ), + } + } + + RETURN_TYPES = ("LATENT","LATENT") + RETURN_NAMES = ("output", "denoised_output") + + FUNCTION = "sample" + + CATEGORY = "sampling/custom_sampling" + + def sample(self, noise, guider, sampler, sigmas, latent_image): + latent = latent_image + latent_image = latent["samples"] + + noise_mask = None + if "noise_mask" in latent: + noise_mask = latent["noise_mask"] + + x0_output = {} + callback = latent_preview.prepare_callback(guider.model_patcher, sigmas.shape[-1] - 1, x0_output) + + disable_pbar = not comfy.utils.PROGRESS_BAR_ENABLED + samples = guider.sample(noise.generate_noise(latent), latent_image, sampler, sigmas, denoise_mask=noise_mask, callback=callback, disable_pbar=disable_pbar, seed=noise.seed) + samples = samples.to(comfy.model_management.intermediate_device()) + + out = latent.copy() + out["samples"] = samples + if "x0" in x0_output: + out_denoised = latent.copy() + out_denoised["samples"] = guider.model_patcher.model.process_latent_out(x0_output["x0"].cpu()) + else: + out_denoised = out + return (out, out_denoised) + NODE_CLASS_MAPPINGS = { "SamplerCustom": SamplerCustom, "BasicScheduler": BasicScheduler, @@ -378,4 +489,9 @@ NODE_CLASS_MAPPINGS = { "SamplerDPMAdaptative": SamplerDPMAdaptative, "SplitSigmas": SplitSigmas, "FlipSigmas": FlipSigmas, + + "CFGGuider": CFGGuider, + "RandomNoise": RandomNoise, + "DisableNoise": DisableNoise, + "SamplerCustomAdvanced": SamplerCustomAdvanced, } From fcfd2bdf8a4fcc57ed842685bb3bfd39b6505e42 Mon Sep 17 00:00:00 2001 From: comfyanonymous Date: Thu, 4 Apr 2024 11:16:49 -0400 Subject: [PATCH 10/29] Small cleanup. --- comfy/samplers.py | 11 +++++++---- comfy_extras/nodes_custom_sampler.py | 2 +- 2 files changed, 8 insertions(+), 5 deletions(-) diff --git a/comfy/samplers.py b/comfy/samplers.py index 57f5632c5..a89e3a6c0 100644 --- a/comfy/samplers.py +++ b/comfy/samplers.py @@ -596,13 +596,16 @@ class CFGGuider: self.original_conds = {} self.cfg = 1.0 - def set_conds(self, conds): - for k in conds: - self.original_conds[k] = comfy.sampler_helpers.convert_cond(conds[k]) + def set_conds(self, positive, negative): + self.inner_set_conds({"positive": positive, "negative": negative}) def set_cfg(self, cfg): self.cfg = cfg + def inner_set_conds(self, conds): + for k in conds: + self.original_conds[k] = comfy.sampler_helpers.convert_cond(conds[k]) + def __call__(self, *args, **kwargs): return self.predict_noise(*args, **kwargs) @@ -646,7 +649,7 @@ class CFGGuider: def sample(model, noise, positive, negative, cfg, device, sampler, sigmas, model_options={}, latent_image=None, denoise_mask=None, callback=None, disable_pbar=False, seed=None): cfg_guider = CFGGuider(model) - cfg_guider.set_conds({"positive": positive, "negative": negative}) + cfg_guider.set_conds(positive, negative) cfg_guider.set_cfg(cfg) return cfg_guider.sample(noise, latent_image, sampler, sigmas, denoise_mask, callback, disable_pbar, seed) diff --git a/comfy_extras/nodes_custom_sampler.py b/comfy_extras/nodes_custom_sampler.py index 32b2a56d6..e9dc3bd98 100644 --- a/comfy_extras/nodes_custom_sampler.py +++ b/comfy_extras/nodes_custom_sampler.py @@ -397,7 +397,7 @@ class CFGGuider: def get_guider(self, model, positive, negative, cfg): guider = comfy.samplers.CFGGuider(model) - guider.set_conds({"positive": positive, "negative": negative}) + guider.set_conds(positive, negative) guider.set_cfg(cfg) return (guider,) From c6bd456c45fd24818223bd4f61a6840e281ac82f Mon Sep 17 00:00:00 2001 From: comfyanonymous Date: Thu, 4 Apr 2024 11:38:25 -0400 Subject: [PATCH 11/29] Make zero denoise a NOP. --- comfy/samplers.py | 12 +++++++++--- comfy_extras/nodes_custom_sampler.py | 5 +++++ 2 files changed, 14 insertions(+), 3 deletions(-) diff --git a/comfy/samplers.py b/comfy/samplers.py index a89e3a6c0..475b1aad6 100644 --- a/comfy/samplers.py +++ b/comfy/samplers.py @@ -624,6 +624,9 @@ class CFGGuider: return self.inner_model.process_latent_out(samples.to(torch.float32)) def sample(self, noise, latent_image, sampler, sigmas, denoise_mask=None, callback=None, disable_pbar=False, seed=None): + if sigmas.shape[-1] == 0: + return latent_image + self.conds = {} for k in self.original_conds: self.conds[k] = list(map(lambda a: a.copy(), self.original_conds[k])) @@ -722,9 +725,12 @@ class KSampler: if denoise is None or denoise > 0.9999: self.sigmas = self.calculate_sigmas(steps).to(self.device) else: - new_steps = int(steps/denoise) - sigmas = self.calculate_sigmas(new_steps).to(self.device) - self.sigmas = sigmas[-(steps + 1):] + if denoise <= 0.0: + self.sigmas = torch.FloatTensor([]) + else: + new_steps = int(steps/denoise) + sigmas = self.calculate_sigmas(new_steps).to(self.device) + self.sigmas = sigmas[-(steps + 1):] def sample(self, noise, positive, negative, cfg, latent_image=None, start_step=None, last_step=None, force_full_denoise=False, denoise_mask=None, sigmas=None, callback=None, disable_pbar=False, seed=None): if sigmas is None: diff --git a/comfy_extras/nodes_custom_sampler.py b/comfy_extras/nodes_custom_sampler.py index e9dc3bd98..a99dbcee8 100644 --- a/comfy_extras/nodes_custom_sampler.py +++ b/comfy_extras/nodes_custom_sampler.py @@ -24,6 +24,8 @@ class BasicScheduler: def get_sigmas(self, model, scheduler, steps, denoise): total_steps = steps if denoise < 1.0: + if denoise <= 0.0: + return (torch.FloatTensor([]),) total_steps = int(steps/denoise) comfy.model_management.load_models_gpu([model]) @@ -160,6 +162,9 @@ class FlipSigmas: FUNCTION = "get_sigmas" def get_sigmas(self, sigmas): + if len(sigmas) == 0: + return (sigmas,) + sigmas = sigmas.flip(0) if sigmas[0] == 0: sigmas[0] = 0.0001 From cfbf3be54b4730d0a434115fcad2ff138a5753ba Mon Sep 17 00:00:00 2001 From: comfyanonymous Date: Thu, 4 Apr 2024 13:57:32 -0400 Subject: [PATCH 12/29] Add basic guider for models with no cfg. --- comfy_extras/nodes_custom_sampler.py | 23 +++++++++++++++++++++++ 1 file changed, 23 insertions(+) diff --git a/comfy_extras/nodes_custom_sampler.py b/comfy_extras/nodes_custom_sampler.py index a99dbcee8..5da437a6c 100644 --- a/comfy_extras/nodes_custom_sampler.py +++ b/comfy_extras/nodes_custom_sampler.py @@ -383,6 +383,28 @@ class SamplerCustom: out_denoised = out return (out, out_denoised) +class Guider_Basic(comfy.samplers.CFGGuider): + def set_conds(self, positive): + self.inner_set_conds({"positive": positive}) + +class BasicGuider: + @classmethod + def INPUT_TYPES(s): + return {"required": + {"model": ("MODEL",), + "conditioning": ("CONDITIONING", ), + } + } + + RETURN_TYPES = ("GUIDER",) + + FUNCTION = "get_guider" + CATEGORY = "sampling/custom_sampling/guiders" + + def get_guider(self, model, conditioning): + guider = Guider_Basic(model) + guider.set_conds(conditioning) + return (guider,) class CFGGuider: @classmethod @@ -496,6 +518,7 @@ NODE_CLASS_MAPPINGS = { "FlipSigmas": FlipSigmas, "CFGGuider": CFGGuider, + "BasicGuider": BasicGuider, "RandomNoise": RandomNoise, "DisableNoise": DisableNoise, "SamplerCustomAdvanced": SamplerCustomAdvanced, From 5272fd4b0389c6e702493d193a1f824f9fa4c7b8 Mon Sep 17 00:00:00 2001 From: comfyanonymous Date: Thu, 4 Apr 2024 14:57:44 -0400 Subject: [PATCH 13/29] Add DualCFGGuider used in IP2P models for example. --- comfy_extras/nodes_custom_sampler.py | 36 ++++++++++++++++++++++++++++ 1 file changed, 36 insertions(+) diff --git a/comfy_extras/nodes_custom_sampler.py b/comfy_extras/nodes_custom_sampler.py index 5da437a6c..fbd8cd25a 100644 --- a/comfy_extras/nodes_custom_sampler.py +++ b/comfy_extras/nodes_custom_sampler.py @@ -428,6 +428,41 @@ class CFGGuider: guider.set_cfg(cfg) return (guider,) +class Guider_DualCFG(comfy.samplers.CFGGuider): + def set_cfg(self, cfg1, cfg2): + self.cfg1 = cfg1 + self.cfg2 = cfg2 + + def set_conds(self, positive, middle, negative): + self.inner_set_conds({"positive": positive, "middle": middle, "negative": negative}) + + def predict_noise(self, x, timestep, model_options={}, seed=None): + out = comfy.samplers.calc_cond_batch(self.inner_model, [self.conds.get("negative", None), self.conds.get("middle", None), self.conds.get("positive", None)], x, timestep, model_options) + return comfy.samplers.cfg_function(self.inner_model, out[1], out[0], self.cfg2, x, timestep, model_options=model_options) + (out[2] - out[1]) * self.cfg1 + +class DualCFGGuider: + @classmethod + def INPUT_TYPES(s): + return {"required": + {"model": ("MODEL",), + "cond1": ("CONDITIONING", ), + "cond2": ("CONDITIONING", ), + "negative": ("CONDITIONING", ), + "cfg_conds": ("FLOAT", {"default": 8.0, "min": 0.0, "max": 100.0, "step":0.1, "round": 0.01}), + "cfg_cond2_negative": ("FLOAT", {"default": 8.0, "min": 0.0, "max": 100.0, "step":0.1, "round": 0.01}), + } + } + + RETURN_TYPES = ("GUIDER",) + + FUNCTION = "get_guider" + CATEGORY = "sampling/custom_sampling/guiders" + + def get_guider(self, model, cond1, cond2, negative, cfg_conds, cfg_cond2_negative): + guider = Guider_DualCFG(model) + guider.set_conds(cond1, cond2, negative) + guider.set_cfg(cfg_conds, cfg_cond2_negative) + return (guider,) class DisableNoise: @classmethod @@ -518,6 +553,7 @@ NODE_CLASS_MAPPINGS = { "FlipSigmas": FlipSigmas, "CFGGuider": CFGGuider, + "DualCFGGuider": DualCFGGuider, "BasicGuider": BasicGuider, "RandomNoise": RandomNoise, "DisableNoise": DisableNoise, From 1f8d8e6c772fe462f697486626c6c1720b16a37d Mon Sep 17 00:00:00 2001 From: comfyanonymous Date: Thu, 4 Apr 2024 15:06:17 -0400 Subject: [PATCH 14/29] Add InstructPixToPixConditioning node. --- comfy_extras/nodes_ip2p.py | 45 ++++++++++++++++++++++++++++++++++++++ nodes.py | 1 + 2 files changed, 46 insertions(+) create mode 100644 comfy_extras/nodes_ip2p.py diff --git a/comfy_extras/nodes_ip2p.py b/comfy_extras/nodes_ip2p.py new file mode 100644 index 000000000..c2e70a84c --- /dev/null +++ b/comfy_extras/nodes_ip2p.py @@ -0,0 +1,45 @@ +import torch + +class InstructPixToPixConditioning: + @classmethod + def INPUT_TYPES(s): + return {"required": {"positive": ("CONDITIONING", ), + "negative": ("CONDITIONING", ), + "vae": ("VAE", ), + "pixels": ("IMAGE", ), + }} + + RETURN_TYPES = ("CONDITIONING","CONDITIONING","LATENT") + RETURN_NAMES = ("positive", "negative", "latent") + FUNCTION = "encode" + + CATEGORY = "conditioning/instructpix2pix" + + def encode(self, positive, negative, pixels, vae): + x = (pixels.shape[1] // 8) * 8 + y = (pixels.shape[2] // 8) * 8 + + if pixels.shape[1] != x or pixels.shape[2] != y: + x_offset = (pixels.shape[1] % 8) // 2 + y_offset = (pixels.shape[2] % 8) // 2 + pixels = pixels[:,x_offset:x + x_offset, y_offset:y + y_offset,:] + + concat_latent = vae.encode(pixels) + + out_latent = {} + out_latent["samples"] = torch.zeros_like(concat_latent) + + out = [] + for conditioning in [positive, negative]: + c = [] + for t in conditioning: + d = t[1].copy() + d["concat_latent_image"] = concat_latent + n = [t[0], d] + c.append(n) + out.append(c) + return (out[0], out[1], out_latent) + +NODE_CLASS_MAPPINGS = { + "InstructPixToPixConditioning": InstructPixToPixConditioning, +} diff --git a/nodes.py b/nodes.py index 5cf020df2..cdb82f801 100644 --- a/nodes.py +++ b/nodes.py @@ -1965,6 +1965,7 @@ def init_custom_nodes(): "nodes_morphology.py", "nodes_stable_cascade.py", "nodes_differential_diffusion.py", + "nodes_ip2p.py", ] import_failed = [] From 1a0486bb96fb1ff10f4ea3c0d62eb815e9630585 Mon Sep 17 00:00:00 2001 From: comfyanonymous Date: Thu, 4 Apr 2024 22:08:49 -0400 Subject: [PATCH 15/29] Fix model needing to be loaded on GPU to generate the sigmas. --- comfy/model_patcher.py | 6 ++++++ comfy/samplers.py | 28 ++++++++++++++-------------- comfy_extras/nodes_custom_sampler.py | 3 +-- 3 files changed, 21 insertions(+), 16 deletions(-) diff --git a/comfy/model_patcher.py b/comfy/model_patcher.py index 8dda84cfd..97fabd4f2 100644 --- a/comfy/model_patcher.py +++ b/comfy/model_patcher.py @@ -150,6 +150,12 @@ class ModelPatcher: def add_object_patch(self, name, obj): self.object_patches[name] = obj + def get_model_object(self, name): + if name in self.object_patches: + return self.object_patches[name] + else: + return comfy.utils.get_attr(self.model, name) + def model_patches_to(self, device): to = self.model_options["transformer_options"] if "patches" in to: diff --git a/comfy/samplers.py b/comfy/samplers.py index 475b1aad6..08cbab756 100644 --- a/comfy/samplers.py +++ b/comfy/samplers.py @@ -274,8 +274,8 @@ class KSamplerX0Inpaint: out = out * denoise_mask + self.latent_image * latent_mask return out -def simple_scheduler(model, steps): - s = model.model_sampling +def simple_scheduler(model_sampling, steps): + s = model_sampling sigs = [] ss = len(s.sigmas) / steps for x in range(steps): @@ -283,8 +283,8 @@ def simple_scheduler(model, steps): sigs += [0.0] return torch.FloatTensor(sigs) -def ddim_scheduler(model, steps): - s = model.model_sampling +def ddim_scheduler(model_sampling, steps): + s = model_sampling sigs = [] ss = max(len(s.sigmas) // steps, 1) x = 1 @@ -295,8 +295,8 @@ def ddim_scheduler(model, steps): sigs += [0.0] return torch.FloatTensor(sigs) -def normal_scheduler(model, steps, sgm=False, floor=False): - s = model.model_sampling +def normal_scheduler(model_sampling, steps, sgm=False, floor=False): + s = model_sampling start = s.timestep(s.sigma_max) end = s.timestep(s.sigma_min) @@ -660,19 +660,19 @@ def sample(model, noise, positive, negative, cfg, device, sampler, sigmas, model SCHEDULER_NAMES = ["normal", "karras", "exponential", "sgm_uniform", "simple", "ddim_uniform"] SAMPLER_NAMES = KSAMPLER_NAMES + ["ddim", "uni_pc", "uni_pc_bh2"] -def calculate_sigmas_scheduler(model, scheduler_name, steps): +def calculate_sigmas(model_sampling, scheduler_name, steps): if scheduler_name == "karras": - sigmas = k_diffusion_sampling.get_sigmas_karras(n=steps, sigma_min=float(model.model_sampling.sigma_min), sigma_max=float(model.model_sampling.sigma_max)) + sigmas = k_diffusion_sampling.get_sigmas_karras(n=steps, sigma_min=float(model_sampling.sigma_min), sigma_max=float(model_sampling.sigma_max)) elif scheduler_name == "exponential": - sigmas = k_diffusion_sampling.get_sigmas_exponential(n=steps, sigma_min=float(model.model_sampling.sigma_min), sigma_max=float(model.model_sampling.sigma_max)) + sigmas = k_diffusion_sampling.get_sigmas_exponential(n=steps, sigma_min=float(model_sampling.sigma_min), sigma_max=float(model_sampling.sigma_max)) elif scheduler_name == "normal": - sigmas = normal_scheduler(model, steps) + sigmas = normal_scheduler(model_sampling, steps) elif scheduler_name == "simple": - sigmas = simple_scheduler(model, steps) + sigmas = simple_scheduler(model_sampling, steps) elif scheduler_name == "ddim_uniform": - sigmas = ddim_scheduler(model, steps) + sigmas = ddim_scheduler(model_sampling, steps) elif scheduler_name == "sgm_uniform": - sigmas = normal_scheduler(model, steps, sgm=True) + sigmas = normal_scheduler(model_sampling, steps, sgm=True) else: logging.error("error invalid scheduler {}".format(scheduler_name)) return sigmas @@ -714,7 +714,7 @@ class KSampler: steps += 1 discard_penultimate_sigma = True - sigmas = calculate_sigmas_scheduler(self.model.model, self.scheduler, steps) + sigmas = calculate_sigmas(self.model.get_model_object("model_sampling"), self.scheduler, steps) if discard_penultimate_sigma: sigmas = torch.cat([sigmas[:-2], sigmas[-1:]]) diff --git a/comfy_extras/nodes_custom_sampler.py b/comfy_extras/nodes_custom_sampler.py index fbd8cd25a..fa1131925 100644 --- a/comfy_extras/nodes_custom_sampler.py +++ b/comfy_extras/nodes_custom_sampler.py @@ -28,8 +28,7 @@ class BasicScheduler: return (torch.FloatTensor([]),) total_steps = int(steps/denoise) - comfy.model_management.load_models_gpu([model]) - sigmas = comfy.samplers.calculate_sigmas_scheduler(model.model, scheduler, total_steps).cpu() + sigmas = comfy.samplers.calculate_sigmas(model.get_model_object("model_sampling"), scheduler, total_steps).cpu() sigmas = sigmas[-(steps + 1):] return (sigmas, ) From 1f4fc9ea0ccceba2e86668a22d86e63b3d262b83 Mon Sep 17 00:00:00 2001 From: comfyanonymous Date: Thu, 4 Apr 2024 23:01:02 -0400 Subject: [PATCH 16/29] Fix issue with get_model_object on patched model. --- comfy/model_patcher.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/comfy/model_patcher.py b/comfy/model_patcher.py index 97fabd4f2..1595d4a2c 100644 --- a/comfy/model_patcher.py +++ b/comfy/model_patcher.py @@ -154,7 +154,10 @@ class ModelPatcher: if name in self.object_patches: return self.object_patches[name] else: - return comfy.utils.get_attr(self.model, name) + if name in self.object_patches_backup: + return self.object_patches_backup[name] + else: + return comfy.utils.get_attr(self.model, name) def model_patches_to(self, device): to = self.model_options["transformer_options"] From 0f5768e038343a8eb96074c8c7f3911abf12461e Mon Sep 17 00:00:00 2001 From: comfyanonymous Date: Thu, 4 Apr 2024 23:38:57 -0400 Subject: [PATCH 17/29] Fix missing arguments in cfg_function. --- comfy/samplers.py | 4 ++-- comfy_extras/nodes_custom_sampler.py | 7 +++++-- 2 files changed, 7 insertions(+), 4 deletions(-) diff --git a/comfy/samplers.py b/comfy/samplers.py index 08cbab756..415a35cc3 100644 --- a/comfy/samplers.py +++ b/comfy/samplers.py @@ -231,7 +231,7 @@ def calc_cond_uncond_batch(model, cond, uncond, x_in, timestep, model_options): logging.warning("WARNING: The comfy.samplers.calc_cond_uncond_batch function is deprecated please use the calc_cond_batch one instead.") return tuple(calc_cond_batch(model, [cond, uncond], x_in, timestep, model_options)) -def cfg_function(model, cond_pred, uncond_pred, cond_scale, x, timestep, model_options={}): +def cfg_function(model, cond_pred, uncond_pred, cond_scale, x, timestep, model_options={}, cond=None, uncond=None): if "sampler_cfg_function" in model_options: args = {"cond": x - cond_pred, "uncond": x - uncond_pred, "cond_scale": cond_scale, "timestep": timestep, "input": x, "sigma": timestep, "cond_denoised": cond_pred, "uncond_denoised": uncond_pred, "model": model, "model_options": model_options} @@ -256,7 +256,7 @@ def sampling_function(model, x, timestep, uncond, cond, cond_scale, model_option conds = [cond, uncond_] out = calc_cond_batch(model, conds, x, timestep, model_options) - return cfg_function(model, out[0], out[1], cond_scale, x, timestep, model_options=model_options) + return cfg_function(model, out[0], out[1], cond_scale, x, timestep, model_options=model_options, cond=cond, uncond=uncond_) class KSamplerX0Inpaint: diff --git a/comfy_extras/nodes_custom_sampler.py b/comfy_extras/nodes_custom_sampler.py index fa1131925..a5d791d9c 100644 --- a/comfy_extras/nodes_custom_sampler.py +++ b/comfy_extras/nodes_custom_sampler.py @@ -436,8 +436,11 @@ class Guider_DualCFG(comfy.samplers.CFGGuider): self.inner_set_conds({"positive": positive, "middle": middle, "negative": negative}) def predict_noise(self, x, timestep, model_options={}, seed=None): - out = comfy.samplers.calc_cond_batch(self.inner_model, [self.conds.get("negative", None), self.conds.get("middle", None), self.conds.get("positive", None)], x, timestep, model_options) - return comfy.samplers.cfg_function(self.inner_model, out[1], out[0], self.cfg2, x, timestep, model_options=model_options) + (out[2] - out[1]) * self.cfg1 + negative_cond = self.conds.get("negative", None) + middle_cond = self.conds.get("middle", None) + + out = comfy.samplers.calc_cond_batch(self.inner_model, [negative_cond, middle_cond, self.conds.get("positive", None)], x, timestep, model_options) + return comfy.samplers.cfg_function(self.inner_model, out[1], out[0], self.cfg2, x, timestep, model_options=model_options, cond=middle_cond, uncond=negative_cond) + (out[2] - out[1]) * self.cfg1 class DualCFGGuider: @classmethod From 41ed7e85ea4bb9e77f44bc9b2f87cca4e3dc99b2 Mon Sep 17 00:00:00 2001 From: comfyanonymous Date: Fri, 5 Apr 2024 00:22:44 -0400 Subject: [PATCH 18/29] Fix object_patches_backup not being the same object across clones. --- comfy/model_patcher.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/comfy/model_patcher.py b/comfy/model_patcher.py index 1595d4a2c..76485ced1 100644 --- a/comfy/model_patcher.py +++ b/comfy/model_patcher.py @@ -471,4 +471,4 @@ class ModelPatcher: for k in keys: comfy.utils.set_attr(self.model, k, self.object_patches_backup[k]) - self.object_patches_backup = {} + self.object_patches_backup.clear() From 1088d1850f9b13233e3cf4460ee077b15e4f712f Mon Sep 17 00:00:00 2001 From: comfyanonymous Date: Fri, 5 Apr 2024 10:40:27 -0400 Subject: [PATCH 19/29] Support for CosXL models. --- comfy/model_base.py | 6 ++++-- comfy/supported_models.py | 7 ++++++- 2 files changed, 10 insertions(+), 3 deletions(-) diff --git a/comfy/model_base.py b/comfy/model_base.py index 898adb665..8c89adf5e 100644 --- a/comfy/model_base.py +++ b/comfy/model_base.py @@ -503,8 +503,10 @@ class SD15_instructpix2pix(IP2P, BaseModel): class SDXL_instructpix2pix(IP2P, SDXL): def __init__(self, model_config, model_type=ModelType.EPS, device=None): super().__init__(model_config, model_type, device=device) - # self.process_ip2p_image_in = lambda image: comfy.latent_formats.SDXL().process_in(image) - self.process_ip2p_image_in = lambda image: image + if model_type == ModelType.V_PREDICTION_EDM: + self.process_ip2p_image_in = lambda image: comfy.latent_formats.SDXL().process_in(image) #cosxl ip2p + else: + self.process_ip2p_image_in = lambda image: image #diffusers ip2p class StableCascade_C(BaseModel): diff --git a/comfy/supported_models.py b/comfy/supported_models.py index 9bfe3ea1d..b3b69e05b 100644 --- a/comfy/supported_models.py +++ b/comfy/supported_models.py @@ -174,6 +174,11 @@ class SDXL(supported_models_base.BASE): self.sampling_settings["sigma_max"] = 80.0 self.sampling_settings["sigma_min"] = 0.002 return model_base.ModelType.EDM + elif "edm_vpred.sigma_max" in state_dict: + self.sampling_settings["sigma_max"] = float(state_dict["edm_vpred.sigma_max"].item()) + if "edm_vpred.sigma_min" in state_dict: + self.sampling_settings["sigma_min"] = float(state_dict["edm_vpred.sigma_min"].item()) + return model_base.ModelType.V_PREDICTION_EDM elif "v_pred" in state_dict: return model_base.ModelType.V_PREDICTION else: @@ -469,7 +474,7 @@ class SDXL_instructpix2pix(SDXL): } def get_model(self, state_dict, prefix="", device=None): - return model_base.SDXL_instructpix2pix(self, device=device) + return model_base.SDXL_instructpix2pix(self, model_type=self.model_type(state_dict, prefix), device=device) models = [Stable_Zero123, SD15_instructpix2pix, SD15, SD20, SD21UnclipL, SD21UnclipH, SDXL_instructpix2pix, SDXLRefiner, SDXL, SSD1B, KOALA_700M, KOALA_1B, Segmind_Vega, SD_X4Upscaler, Stable_Cascade_C, Stable_Cascade_B, SV3D_u, SV3D_p] From ea9ac9d30beb23119e590610d3ec5dcd146a12f2 Mon Sep 17 00:00:00 2001 From: comfyanonymous Date: Fri, 5 Apr 2024 11:33:36 -0400 Subject: [PATCH 20/29] Fix PerpNeg node. --- comfy_extras/nodes_perpneg.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/comfy_extras/nodes_perpneg.py b/comfy_extras/nodes_perpneg.py index 9e8a218f9..992e8ab62 100644 --- a/comfy_extras/nodes_perpneg.py +++ b/comfy_extras/nodes_perpneg.py @@ -1,10 +1,11 @@ import torch import comfy.model_management -import comfy.sample +import comfy.sampler_helpers import comfy.samplers import comfy.utils +#TODO: This node should be removed and replaced with one that uses the new Guider/SamplerCustomAdvanced. class PerpNeg: @classmethod def INPUT_TYPES(s): @@ -19,7 +20,7 @@ class PerpNeg: def patch(self, model, empty_conditioning, neg_scale): m = model.clone() - nocond = comfy.sample.convert_cond(empty_conditioning) + nocond = comfy.sampler_helpers.convert_cond(empty_conditioning) def cfg_function(args): model = args["model"] From 38ed2da2ddc463a84acc08dbd2bbd5785190d487 Mon Sep 17 00:00:00 2001 From: kk-89 <31396340+kk-89@users.noreply.github.com> Date: Fri, 5 Apr 2024 09:02:13 -0700 Subject: [PATCH 21/29] Fix typo in lowvram patcher (#3209) --- comfy/model_patcher.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/comfy/model_patcher.py b/comfy/model_patcher.py index 76485ced1..657ddf843 100644 --- a/comfy/model_patcher.py +++ b/comfy/model_patcher.py @@ -287,7 +287,7 @@ class ModelPatcher: if weight_key in self.patches: m.weight_function = LowVramPatch(weight_key, self) if bias_key in self.patches: - m.bias_function = LowVramPatch(weight_key, self) + m.bias_function = LowVramPatch(bias_key, self) m.prev_comfy_cast_weights = m.comfy_cast_weights m.comfy_cast_weights = True From a7dd82e668bfaf7fac365a4e73a1ba1acf224fbb Mon Sep 17 00:00:00 2001 From: comfyanonymous Date: Fri, 5 Apr 2024 14:59:05 -0400 Subject: [PATCH 22/29] Fix copy paste issue with litegraph. --- web/lib/litegraph.core.js | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/web/lib/litegraph.core.js b/web/lib/litegraph.core.js index 4ff05ae81..427a62b59 100644 --- a/web/lib/litegraph.core.js +++ b/web/lib/litegraph.core.js @@ -7247,7 +7247,7 @@ LGraphNode.prototype.executeAction = function(action) //create links for (var i = 0; i < clipboard_info.links.length; ++i) { var link_info = clipboard_info.links[i]; - var origin_node; + var origin_node = undefined; var origin_node_relative_id = link_info[0]; if (origin_node_relative_id != null) { origin_node = nodes[origin_node_relative_id]; From d8dea4cdb88cae92ac4b4a4d79bbfd60b7f8923d Mon Sep 17 00:00:00 2001 From: comfyanonymous Date: Fri, 5 Apr 2024 21:36:23 -0400 Subject: [PATCH 23/29] Fix DisableNoise node. --- comfy_extras/nodes_custom_sampler.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/comfy_extras/nodes_custom_sampler.py b/comfy_extras/nodes_custom_sampler.py index a5d791d9c..56cb505e6 100644 --- a/comfy_extras/nodes_custom_sampler.py +++ b/comfy_extras/nodes_custom_sampler.py @@ -320,7 +320,7 @@ class Noise_EmptyNoise: def generate_noise(self, input_latent): latent_image = input_latent["samples"] - return torch.zeros(shape, dtype=latent_image.dtype, layout=latent_image.layout, device="cpu") + return torch.zeros(latent_image.shape, dtype=latent_image.dtype, layout=latent_image.layout, device="cpu") class Noise_RandomNoise: @@ -477,7 +477,7 @@ class DisableNoise: FUNCTION = "get_noise" CATEGORY = "sampling/custom_sampling/noise" - def get_noise(self, noise_seed): + def get_noise(self): return (Noise_EmptyNoise(),) From de172f8be716ab8a28c202d8162e62a184f39acd Mon Sep 17 00:00:00 2001 From: Gorka Eguileor Date: Sat, 6 Apr 2024 18:10:17 +0200 Subject: [PATCH 24/29] Improve A1111 metadata parsing (#3216) * A1111 import: Set VAE name This patch sets the VAE name for the `VAELoader` when present in the png metadata. * A1111 import: Skip all hashes When importing from A1111 the parsing assumes that values of a key will never contain a ":", which is not correct. There are 2 cases where we can have ":" in the value: - Inside a string. E.g.: Lora hashes: "xl_more_art-full_v1: fe3b4816be83, add-detail-xl: 9c783c8ce46c" - When the value is a json dictionary. E.g.: Hashes: {"vae": "63aeecb90f", "embed:negativeXL_D": "fff5d51ab6"} This patch changes how we parse the metadata to take those 2 cases into account and also skips the following additional keys that are present in some Forge images: - Version - VAE hash - TI hashes - Lora hashes - Hashes * A1111 import: Parse Hires steps This patch parses the `Hires steps` parameter that is part of the High Resolution Upscale configuration when it is present, and fallbacks to the one from the `samplerNode` (like the code currently does) if it's not present. --- web/scripts/pnginfo.js | 14 +++++++++++--- 1 file changed, 11 insertions(+), 3 deletions(-) diff --git a/web/scripts/pnginfo.js b/web/scripts/pnginfo.js index 169609209..7132fb60f 100644 --- a/web/scripts/pnginfo.js +++ b/web/scripts/pnginfo.js @@ -170,9 +170,12 @@ export async function importA1111(graph, parameters) { const opts = parameters .substr(p) .split("\n")[1] - .split(",") + .match(new RegExp("\\s*([^:]+:\\s*([^\"\\{].*?|\".*?\"|\\{.*?\\}))\\s*(,|$)", "g")) .reduce((p, n) => { const s = n.split(":"); + if (s[1].endsWith(',')) { + s[1] = s[1].substr(0, s[1].length -1); + } p[s[0].trim().toLowerCase()] = s[1].trim(); return p; }, {}); @@ -191,6 +194,7 @@ export async function importA1111(graph, parameters) { const vaeLoaderNode = LiteGraph.createNode("VAELoader"); const saveNode = LiteGraph.createNode("SaveImage"); let hrSamplerNode = null; + let hrSteps = null; const ceil64 = (v) => Math.ceil(v / 64) * 64; @@ -290,6 +294,9 @@ export async function importA1111(graph, parameters) { model(v) { setWidgetValue(ckptNode, "ckpt_name", v, true); }, + "vae"(v) { + setWidgetValue(vaeLoaderNode, "vae_name", v, true); + }, "cfg scale"(v) { setWidgetValue(samplerNode, "cfg", +v); }, @@ -316,6 +323,7 @@ export async function importA1111(graph, parameters) { const h = ceil64(+wxh[1]); const hrUp = popOpt("hires upscale"); const hrSz = popOpt("hires resize"); + hrSteps = popOpt("hires steps"); let hrMethod = popOpt("hires upscaler"); setWidgetValue(imageNode, "width", w); @@ -398,7 +406,7 @@ export async function importA1111(graph, parameters) { } if (hrSamplerNode) { - setWidgetValue(hrSamplerNode, "steps", getWidget(samplerNode, "steps").value); + setWidgetValue(hrSamplerNode, "steps", hrSteps? +hrSteps : getWidget(samplerNode, "steps").value); setWidgetValue(hrSamplerNode, "cfg", getWidget(samplerNode, "cfg").value); setWidgetValue(hrSamplerNode, "scheduler", getWidget(samplerNode, "scheduler").value); setWidgetValue(hrSamplerNode, "sampler_name", getWidget(samplerNode, "sampler_name").value); @@ -415,7 +423,7 @@ export async function importA1111(graph, parameters) { graph.arrange(); - for (const opt of ["model hash", "ensd"]) { + for (const opt of ["model hash", "ensd", "version", "vae hash", "ti hashes", "lora hashes", "hashes"]) { delete opts[opt]; } From 0a03009808a5ad13fa3a44edbabcae68576c3982 Mon Sep 17 00:00:00 2001 From: comfyanonymous Date: Sat, 6 Apr 2024 18:38:39 -0400 Subject: [PATCH 25/29] Fix issue with controlnet models getting loaded multiple times. --- comfy/controlnet.py | 10 +++++++--- comfy/model_management.py | 2 ++ 2 files changed, 9 insertions(+), 3 deletions(-) diff --git a/comfy/controlnet.py b/comfy/controlnet.py index b6941d8c4..8cf4a61a6 100644 --- a/comfy/controlnet.py +++ b/comfy/controlnet.py @@ -138,11 +138,13 @@ class ControlBase: return out class ControlNet(ControlBase): - def __init__(self, control_model, global_average_pooling=False, device=None, load_device=None, manual_cast_dtype=None): + def __init__(self, control_model=None, global_average_pooling=False, device=None, load_device=None, manual_cast_dtype=None): super().__init__(device) self.control_model = control_model self.load_device = load_device - self.control_model_wrapped = comfy.model_patcher.ModelPatcher(self.control_model, load_device=load_device, offload_device=comfy.model_management.unet_offload_device()) + if control_model is not None: + self.control_model_wrapped = comfy.model_patcher.ModelPatcher(self.control_model, load_device=load_device, offload_device=comfy.model_management.unet_offload_device()) + self.global_average_pooling = global_average_pooling self.model_sampling_current = None self.manual_cast_dtype = manual_cast_dtype @@ -183,7 +185,9 @@ class ControlNet(ControlBase): return self.control_merge(None, control, control_prev, output_dtype) def copy(self): - c = ControlNet(self.control_model, global_average_pooling=self.global_average_pooling, load_device=self.load_device, manual_cast_dtype=self.manual_cast_dtype) + c = ControlNet(None, global_average_pooling=self.global_average_pooling, load_device=self.load_device, manual_cast_dtype=self.manual_cast_dtype) + c.control_model = self.control_model + c.control_model_wrapped = self.control_model_wrapped self.copy_to(c) return c diff --git a/comfy/model_management.py b/comfy/model_management.py index 26216432a..310ec2537 100644 --- a/comfy/model_management.py +++ b/comfy/model_management.py @@ -385,6 +385,8 @@ def load_models_gpu(models, memory_required=0): inference_memory = minimum_inference_memory() extra_mem = max(inference_memory, memory_required) + models = set(models) + models_to_load = [] models_already_loaded = [] for x in models: From 80bda6c16393e7af7934e791b1babedb2cf4896a Mon Sep 17 00:00:00 2001 From: comfyanonymous Date: Sun, 7 Apr 2024 14:27:40 -0400 Subject: [PATCH 26/29] Cleanup a few conditioning nodes. --- node_helpers.py | 10 ++++++++++ nodes.py | 39 ++++++++++++--------------------------- 2 files changed, 22 insertions(+), 27 deletions(-) create mode 100644 node_helpers.py diff --git a/node_helpers.py b/node_helpers.py new file mode 100644 index 000000000..8828a4ec9 --- /dev/null +++ b/node_helpers.py @@ -0,0 +1,10 @@ + +def conditioning_set_values(conditioning, values={}): + c = [] + for t in conditioning: + n = [t[0], t[1].copy()] + for k in values: + n[1][k] = values[k] + c.append(n) + + return c diff --git a/nodes.py b/nodes.py index cdb82f801..6d625cc2c 100644 --- a/nodes.py +++ b/nodes.py @@ -34,6 +34,7 @@ import importlib import folder_paths import latent_preview +import node_helpers def before_node_execution(): comfy.model_management.throw_exception_if_processing_interrupted() @@ -151,13 +152,9 @@ class ConditioningSetArea: CATEGORY = "conditioning" def append(self, conditioning, width, height, x, y, strength): - c = [] - for t in conditioning: - n = [t[0], t[1].copy()] - n[1]['area'] = (height // 8, width // 8, y // 8, x // 8) - n[1]['strength'] = strength - n[1]['set_area_to_bounds'] = False - c.append(n) + c = node_helpers.conditioning_set_values(conditioning, {"area": (height // 8, width // 8, y // 8, x // 8), + "strength": strength, + "set_area_to_bounds": False}) return (c, ) class ConditioningSetAreaPercentage: @@ -176,13 +173,9 @@ class ConditioningSetAreaPercentage: CATEGORY = "conditioning" def append(self, conditioning, width, height, x, y, strength): - c = [] - for t in conditioning: - n = [t[0], t[1].copy()] - n[1]['area'] = ("percentage", height, width, y, x) - n[1]['strength'] = strength - n[1]['set_area_to_bounds'] = False - c.append(n) + c = node_helpers.conditioning_set_values(conditioning, {"area": ("percentage", height, width, y, x), + "strength": strength, + "set_area_to_bounds": False}) return (c, ) class ConditioningSetAreaStrength: @@ -197,11 +190,7 @@ class ConditioningSetAreaStrength: CATEGORY = "conditioning" def append(self, conditioning, strength): - c = [] - for t in conditioning: - n = [t[0], t[1].copy()] - n[1]['strength'] = strength - c.append(n) + c = node_helpers.conditioning_set_values(conditioning, {"strength": strength}) return (c, ) @@ -219,19 +208,15 @@ class ConditioningSetMask: CATEGORY = "conditioning" def append(self, conditioning, mask, set_cond_area, strength): - c = [] set_area_to_bounds = False if set_cond_area != "default": set_area_to_bounds = True if len(mask.shape) < 3: mask = mask.unsqueeze(0) - for t in conditioning: - n = [t[0], t[1].copy()] - _, h, w = mask.shape - n[1]['mask'] = mask - n[1]['set_area_to_bounds'] = set_area_to_bounds - n[1]['mask_strength'] = strength - c.append(n) + + c = node_helpers.conditioning_set_values(conditioning, {"mask": mask, + "set_area_to_bounds": set_area_to_bounds, + "mask_strength": strength}) return (c, ) class ConditioningZeroOut: From c9fc242e2ca7c471743d05648abf7785f0610590 Mon Sep 17 00:00:00 2001 From: comfyanonymous Date: Sun, 7 Apr 2024 14:34:43 -0400 Subject: [PATCH 27/29] The middle prompt should be treated more as a negative prompt. --- comfy_extras/nodes_custom_sampler.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/comfy_extras/nodes_custom_sampler.py b/comfy_extras/nodes_custom_sampler.py index 56cb505e6..1971f2c57 100644 --- a/comfy_extras/nodes_custom_sampler.py +++ b/comfy_extras/nodes_custom_sampler.py @@ -4,6 +4,7 @@ from comfy.k_diffusion import sampling as k_diffusion_sampling import latent_preview import torch import comfy.utils +import node_helpers class BasicScheduler: @@ -433,6 +434,7 @@ class Guider_DualCFG(comfy.samplers.CFGGuider): self.cfg2 = cfg2 def set_conds(self, positive, middle, negative): + middle = node_helpers.conditioning_set_values(middle, {"prompt_type": "negative"}) self.inner_set_conds({"positive": positive, "middle": middle, "negative": negative}) def predict_noise(self, x, timestep, model_options={}, seed=None): From d644b6bcd8991357de674f72eba02d81875d7847 Mon Sep 17 00:00:00 2001 From: comfyanonymous Date: Sun, 7 Apr 2024 14:40:43 -0400 Subject: [PATCH 28/29] Cleanup some more conditioning nodes. --- nodes.py | 18 ++++-------------- 1 file changed, 4 insertions(+), 14 deletions(-) diff --git a/nodes.py b/nodes.py index 6d625cc2c..a1baa98a6 100644 --- a/nodes.py +++ b/nodes.py @@ -251,13 +251,8 @@ class ConditioningSetTimestepRange: CATEGORY = "advanced/conditioning" def set_range(self, conditioning, start, end): - c = [] - for t in conditioning: - d = t[1].copy() - d['start_percent'] = start - d['end_percent'] = end - n = [t[0], d] - c.append(n) + c = node_helpers.conditioning_set_values(conditioning, {"start_percent": start, + "end_percent": end}) return (c, ) class VAEDecode: @@ -398,13 +393,8 @@ class InpaintModelConditioning: out = [] for conditioning in [positive, negative]: - c = [] - for t in conditioning: - d = t[1].copy() - d["concat_latent_image"] = concat_latent - d["concat_mask"] = mask - n = [t[0], d] - c.append(n) + c = node_helpers.conditioning_set_values(conditioning, {"concat_latent_image": concat_latent, + "concat_mask": mask}) out.append(c) return (out[0], out[1], out_latent) From 30abc324c2f73e6b648093ccd4741dece20be1e5 Mon Sep 17 00:00:00 2001 From: comfyanonymous Date: Mon, 8 Apr 2024 00:36:22 -0400 Subject: [PATCH 29/29] Support properly saving CosXL checkpoints. --- comfy/sd.py | 5 ++++- comfy_extras/nodes_model_merging.py | 11 ++++++++++- 2 files changed, 14 insertions(+), 2 deletions(-) diff --git a/comfy/sd.py b/comfy/sd.py index 85821120e..3919f4bfa 100644 --- a/comfy/sd.py +++ b/comfy/sd.py @@ -600,7 +600,7 @@ def load_unet(unet_path): raise RuntimeError("ERROR: Could not detect model type of: {}".format(unet_path)) return model -def save_checkpoint(output_path, model, clip=None, vae=None, clip_vision=None, metadata=None): +def save_checkpoint(output_path, model, clip=None, vae=None, clip_vision=None, metadata=None, extra_keys={}): clip_sd = None load_models = [model] if clip is not None: @@ -610,4 +610,7 @@ def save_checkpoint(output_path, model, clip=None, vae=None, clip_vision=None, m model_management.load_models_gpu(load_models) clip_vision_sd = clip_vision.get_sd() if clip_vision is not None else None sd = model.model.state_dict_for_saving(clip_sd, vae.get_sd(), clip_vision_sd) + for k in extra_keys: + sd[k] = extra_keys[k] + comfy.utils.save_torch_file(sd, output_path, metadata=metadata) diff --git a/comfy_extras/nodes_model_merging.py b/comfy_extras/nodes_model_merging.py index a25b73ca7..2a431f65d 100644 --- a/comfy_extras/nodes_model_merging.py +++ b/comfy_extras/nodes_model_merging.py @@ -2,7 +2,9 @@ import comfy.sd import comfy.utils import comfy.model_base import comfy.model_management +import comfy.model_sampling +import torch import folder_paths import json import os @@ -189,6 +191,13 @@ def save_checkpoint(model, clip=None, vae=None, clip_vision=None, filename_prefi # "stable-diffusion-v2-768-v", "stable-diffusion-v2-unclip-l", "stable-diffusion-v2-unclip-h", # "v2-inpainting" + extra_keys = {} + model_sampling = model.get_model_object("model_sampling") + if isinstance(model_sampling, comfy.model_sampling.ModelSamplingContinuousEDM): + if isinstance(model_sampling, comfy.model_sampling.V_PREDICTION): + extra_keys["edm_vpred.sigma_max"] = torch.tensor(model_sampling.sigma_max).float() + extra_keys["edm_vpred.sigma_min"] = torch.tensor(model_sampling.sigma_min).float() + if model.model.model_type == comfy.model_base.ModelType.EPS: metadata["modelspec.predict_key"] = "epsilon" elif model.model.model_type == comfy.model_base.ModelType.V_PREDICTION: @@ -203,7 +212,7 @@ def save_checkpoint(model, clip=None, vae=None, clip_vision=None, filename_prefi output_checkpoint = f"{filename}_{counter:05}_.safetensors" output_checkpoint = os.path.join(full_output_folder, output_checkpoint) - comfy.sd.save_checkpoint(output_checkpoint, model, clip, vae, clip_vision, metadata=metadata) + comfy.sd.save_checkpoint(output_checkpoint, model, clip, vae, clip_vision, metadata=metadata, extra_keys=extra_keys) class CheckpointSave: def __init__(self):