From f2fe635c9f56a8e78866f59b3f110585e75b42f4 Mon Sep 17 00:00:00 2001 From: comfyanonymous Date: Fri, 15 Mar 2024 19:34:22 -0400 Subject: [PATCH 01/54] SamplerDPMAdaptative node to test the different options. --- comfy/samplers.py | 4 ++-- comfy_extras/nodes_custom_sampler.py | 28 ++++++++++++++++++++++++++++ 2 files changed, 30 insertions(+), 2 deletions(-) diff --git a/comfy/samplers.py b/comfy/samplers.py index 16b4514e1..d721cb2e5 100644 --- a/comfy/samplers.py +++ b/comfy/samplers.py @@ -559,11 +559,11 @@ def ksampler(sampler_name, extra_options={}, inpaint_options={}): return k_diffusion_sampling.sample_dpm_fast(model, noise, sigma_min, sigmas[0], total_steps, extra_args=extra_args, callback=callback, disable=disable) sampler_function = dpm_fast_function elif sampler_name == "dpm_adaptive": - def dpm_adaptive_function(model, noise, sigmas, extra_args, callback, disable): + def dpm_adaptive_function(model, noise, sigmas, extra_args, callback, disable, **extra_options): sigma_min = sigmas[-1] if sigma_min == 0: sigma_min = sigmas[-2] - return k_diffusion_sampling.sample_dpm_adaptive(model, noise, sigma_min, sigmas[0], extra_args=extra_args, callback=callback, disable=disable) + return k_diffusion_sampling.sample_dpm_adaptive(model, noise, sigma_min, sigmas[0], extra_args=extra_args, callback=callback, disable=disable, **extra_options) sampler_function = dpm_adaptive_function else: sampler_function = getattr(k_diffusion_sampling, "sample_{}".format(sampler_name)) diff --git a/comfy_extras/nodes_custom_sampler.py b/comfy_extras/nodes_custom_sampler.py index 8b808ce0b..72ff7957f 100644 --- a/comfy_extras/nodes_custom_sampler.py +++ b/comfy_extras/nodes_custom_sampler.py @@ -283,6 +283,33 @@ class SamplerLMS: sampler = comfy.samplers.ksampler("lms", {"order": order}) return (sampler, ) +class SamplerDPMAdaptative: + @classmethod + def INPUT_TYPES(s): + return {"required": + {"order": ("INT", {"default": 3, "min": 2, "max": 3}), + "rtol": ("FLOAT", {"default": 0.05, "min": 0.0, "max": 100.0, "step":0.01, "round": False}), + "atol": ("FLOAT", {"default": 0.0078, "min": 0.0, "max": 100.0, "step":0.01, "round": False}), + "h_init": ("FLOAT", {"default": 0.05, "min": 0.0, "max": 100.0, "step":0.01, "round": False}), + "pcoeff": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 100.0, "step":0.01, "round": False}), + "icoeff": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 100.0, "step":0.01, "round": False}), + "dcoeff": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 100.0, "step":0.01, "round": False}), + "accept_safety": ("FLOAT", {"default": 0.81, "min": 0.0, "max": 100.0, "step":0.01, "round": False}), + "eta": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 100.0, "step":0.01, "round": False}), + "s_noise": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 100.0, "step":0.01, "round": False}), + } + } + RETURN_TYPES = ("SAMPLER",) + CATEGORY = "sampling/custom_sampling/samplers" + + FUNCTION = "get_sampler" + + def get_sampler(self, order, rtol, atol, h_init, pcoeff, icoeff, dcoeff, accept_safety, eta, s_noise): + sampler = comfy.samplers.ksampler("dpm_adaptive", {"order": order, "rtol": rtol, "atol": atol, "h_init": h_init, "pcoeff": pcoeff, + "icoeff": icoeff, "dcoeff": dcoeff, "accept_safety": accept_safety, "eta": eta, + "s_noise":s_noise }) + return (sampler, ) + class SamplerCustom: @classmethod def INPUT_TYPES(s): @@ -348,6 +375,7 @@ NODE_CLASS_MAPPINGS = { "SamplerDPMPP_3M_SDE": SamplerDPMPP_3M_SDE, "SamplerDPMPP_2M_SDE": SamplerDPMPP_2M_SDE, "SamplerDPMPP_SDE": SamplerDPMPP_SDE, + "SamplerDPMAdaptative": SamplerDPMAdaptative, "SplitSigmas": SplitSigmas, "FlipSigmas": FlipSigmas, } From d7897fff2cfd38eac051fbc958a6f944bdf68cc9 Mon Sep 17 00:00:00 2001 From: comfyanonymous Date: Sat, 16 Mar 2024 14:49:35 -0400 Subject: [PATCH 02/54] Move cascade scale factor from stage_a to latent_formats.py --- comfy/latent_formats.py | 2 +- comfy/ldm/cascade/stage_a.py | 9 +++------ 2 files changed, 4 insertions(+), 7 deletions(-) diff --git a/comfy/latent_formats.py b/comfy/latent_formats.py index 674364e72..4ca466d9a 100644 --- a/comfy/latent_formats.py +++ b/comfy/latent_formats.py @@ -95,7 +95,7 @@ class SC_Prior(LatentFormat): class SC_B(LatentFormat): def __init__(self): - self.scale_factor = 1.0 + self.scale_factor = 1.0 / 0.43 self.latent_rgb_factors = [ [ 0.1121, 0.2006, 0.1023], [-0.2093, -0.0222, -0.0195], diff --git a/comfy/ldm/cascade/stage_a.py b/comfy/ldm/cascade/stage_a.py index 260ccfc0b..ca8867eaf 100644 --- a/comfy/ldm/cascade/stage_a.py +++ b/comfy/ldm/cascade/stage_a.py @@ -163,11 +163,9 @@ class ResBlock(nn.Module): class StageA(nn.Module): - def __init__(self, levels=2, bottleneck_blocks=12, c_hidden=384, c_latent=4, codebook_size=8192, - scale_factor=0.43): # 0.3764 + def __init__(self, levels=2, bottleneck_blocks=12, c_hidden=384, c_latent=4, codebook_size=8192): super().__init__() self.c_latent = c_latent - self.scale_factor = scale_factor c_levels = [c_hidden // (2 ** i) for i in reversed(range(levels))] # Encoder blocks @@ -214,12 +212,11 @@ class StageA(nn.Module): x = self.down_blocks(x) if quantize: qe, (vq_loss, commit_loss), indices = self.vquantizer.forward(x, dim=1) - return qe / self.scale_factor, x / self.scale_factor, indices, vq_loss + commit_loss * 0.25 + return qe, x, indices, vq_loss + commit_loss * 0.25 else: - return x / self.scale_factor + return x def decode(self, x): - x = x * self.scale_factor x = self.up_blocks(x) x = self.out_block(x) return x From d3406d8d588d6b0c2da44c9bf378733a2077a14f Mon Sep 17 00:00:00 2001 From: comfyanonymous Date: Sun, 17 Mar 2024 08:57:49 -0400 Subject: [PATCH 03/54] Increase image batch nodes maximum values. --- comfy_extras/nodes_images.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/comfy_extras/nodes_images.py b/comfy_extras/nodes_images.py index 8f638bf8f..af37666b2 100644 --- a/comfy_extras/nodes_images.py +++ b/comfy_extras/nodes_images.py @@ -37,7 +37,7 @@ class RepeatImageBatch: @classmethod def INPUT_TYPES(s): return {"required": { "image": ("IMAGE",), - "amount": ("INT", {"default": 1, "min": 1, "max": 64}), + "amount": ("INT", {"default": 1, "min": 1, "max": 4096}), }} RETURN_TYPES = ("IMAGE",) FUNCTION = "repeat" @@ -52,8 +52,8 @@ class ImageFromBatch: @classmethod def INPUT_TYPES(s): return {"required": { "image": ("IMAGE",), - "batch_index": ("INT", {"default": 0, "min": 0, "max": 63}), - "length": ("INT", {"default": 1, "min": 1, "max": 64}), + "batch_index": ("INT", {"default": 0, "min": 0, "max": 4095}), + "length": ("INT", {"default": 1, "min": 1, "max": 4096}), }} RETURN_TYPES = ("IMAGE",) FUNCTION = "frombatch" From cacb022c4a5b9614f96086a866c8a4c4e9e85760 Mon Sep 17 00:00:00 2001 From: comfyanonymous Date: Mon, 18 Mar 2024 00:26:23 -0400 Subject: [PATCH 04/54] Make saved SD1 checkpoints match more closely the official one. --- comfy/supported_models.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/comfy/supported_models.py b/comfy/supported_models.py index 375821032..e12935a27 100644 --- a/comfy/supported_models.py +++ b/comfy/supported_models.py @@ -45,6 +45,11 @@ class SD15(supported_models_base.BASE): return state_dict def process_clip_state_dict_for_saving(self, state_dict): + pop_keys = ["clip_l.transformer.text_projection.weight", "clip_l.logit_scale"] + for p in pop_keys: + if p in state_dict: + state_dict.pop(p) + replace_prefix = {"clip_l.": "cond_stage_model."} return utils.state_dict_prefix_replace(state_dict, replace_prefix) From b1a16d4500b89f4c7db2aadf41acb378efd20948 Mon Sep 17 00:00:00 2001 From: comfyanonymous Date: Mon, 18 Mar 2024 13:50:11 -0400 Subject: [PATCH 05/54] Fix stable cascade img2img not working with all resolutions. --- comfy_extras/nodes_stable_cascade.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/comfy_extras/nodes_stable_cascade.py b/comfy_extras/nodes_stable_cascade.py index 7e2d37d22..fcbbeb27f 100644 --- a/comfy_extras/nodes_stable_cascade.py +++ b/comfy_extras/nodes_stable_cascade.py @@ -74,7 +74,7 @@ class StableCascade_StageC_VAEEncode: s = comfy.utils.common_upscale(image.movedim(-1,1), out_width, out_height, "bicubic", "center").movedim(1,-1) c_latent = vae.encode(s[:,:,:,:3]) - b_latent = torch.zeros([c_latent.shape[0], 4, height // 4, width // 4]) + b_latent = torch.zeros([c_latent.shape[0], 4, (height // 8) * 2, (width // 8) * 2]) return ({ "samples": c_latent, }, { From 0b78213bdaaa9021daa870c544be6ce86f54d30d Mon Sep 17 00:00:00 2001 From: comfyanonymous Date: Mon, 18 Mar 2024 15:51:23 -0400 Subject: [PATCH 06/54] Fix neg scale step. --- comfy_extras/nodes_perpneg.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/comfy_extras/nodes_perpneg.py b/comfy_extras/nodes_perpneg.py index 64bbc1dcd..dc73c5528 100644 --- a/comfy_extras/nodes_perpneg.py +++ b/comfy_extras/nodes_perpneg.py @@ -10,7 +10,7 @@ class PerpNeg: def INPUT_TYPES(s): return {"required": {"model": ("MODEL", ), "empty_conditioning": ("CONDITIONING", ), - "neg_scale": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 100.0}), + "neg_scale": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 100.0, "step": 0.01}), }} RETURN_TYPES = ("MODEL",) FUNCTION = "patch" From 40e124c6be01195eada95e8c319ca6ddf4fd1a17 Mon Sep 17 00:00:00 2001 From: comfyanonymous Date: Mon, 18 Mar 2024 10:04:51 -0400 Subject: [PATCH 07/54] SV3D support. --- comfy/model_base.py | 30 +++++++++++++++++ comfy/supported_models.py | 37 ++++++++++++++++++++- comfy_extras/nodes_stable3d.py | 53 +++++++++++++++++++++++++++---- comfy_extras/nodes_video_model.py | 28 ++++++++++++++++ 4 files changed, 141 insertions(+), 7 deletions(-) diff --git a/comfy/model_base.py b/comfy/model_base.py index 5da71e632..bc019de53 100644 --- a/comfy/model_base.py +++ b/comfy/model_base.py @@ -380,6 +380,36 @@ class SVD_img2vid(BaseModel): out['num_video_frames'] = comfy.conds.CONDConstant(noise.shape[0]) return out +class SV3D_u(SVD_img2vid): + def encode_adm(self, **kwargs): + augmentation = kwargs.get("augmentation_level", 0) + + out = [] + out.append(self.embedder(torch.flatten(torch.Tensor([augmentation])))) + + flat = torch.flatten(torch.cat(out)).unsqueeze(dim=0) + return flat + +class SV3D_p(SVD_img2vid): + def __init__(self, model_config, model_type=ModelType.V_PREDICTION_EDM, device=None): + super().__init__(model_config, model_type, device=device) + self.embedder_512 = Timestep(512) + + def encode_adm(self, **kwargs): + augmentation = kwargs.get("augmentation_level", 0) + elevation = kwargs.get("elevation", 0) #elevation and azimuth are in degrees here + azimuth = kwargs.get("azimuth", 0) + noise = kwargs.get("noise", None) + + out = [] + out.append(self.embedder(torch.flatten(torch.Tensor([augmentation])))) + out.append(self.embedder_512(torch.deg2rad(torch.fmod(torch.flatten(90 - torch.Tensor([elevation])), 360.0)))) + out.append(self.embedder_512(torch.deg2rad(torch.fmod(torch.flatten(torch.Tensor([azimuth])), 360.0)))) + + out = list(map(lambda a: utils.resize_to_batch_size(a, noise.shape[0]), out)) + return torch.cat(out, dim=1) + + class Stable_Zero123(BaseModel): def __init__(self, model_config, model_type=ModelType.EPS, device=None, cc_projection_weight=None, cc_projection_bias=None): super().__init__(model_config, model_type, device=device) diff --git a/comfy/supported_models.py b/comfy/supported_models.py index e12935a27..2ce9736b7 100644 --- a/comfy/supported_models.py +++ b/comfy/supported_models.py @@ -284,6 +284,41 @@ class SVD_img2vid(supported_models_base.BASE): def clip_target(self): return None +class SV3D_u(SVD_img2vid): + unet_config = { + "model_channels": 320, + "in_channels": 8, + "use_linear_in_transformer": True, + "transformer_depth": [1, 1, 1, 1, 1, 1, 0, 0], + "context_dim": 1024, + "adm_in_channels": 256, + "use_temporal_attention": True, + "use_temporal_resblock": True + } + + vae_key_prefix = ["conditioner.embedders.1.encoder."] + + def get_model(self, state_dict, prefix="", device=None): + out = model_base.SV3D_u(self, device=device) + return out + +class SV3D_p(SV3D_u): + unet_config = { + "model_channels": 320, + "in_channels": 8, + "use_linear_in_transformer": True, + "transformer_depth": [1, 1, 1, 1, 1, 1, 0, 0], + "context_dim": 1024, + "adm_in_channels": 1280, + "use_temporal_attention": True, + "use_temporal_resblock": True + } + + + def get_model(self, state_dict, prefix="", device=None): + out = model_base.SV3D_p(self, device=device) + return out + class Stable_Zero123(supported_models_base.BASE): unet_config = { "context_dim": 768, @@ -405,5 +440,5 @@ class Stable_Cascade_B(Stable_Cascade_C): return out -models = [Stable_Zero123, SD15, SD20, SD21UnclipL, SD21UnclipH, SDXLRefiner, SDXL, SSD1B, KOALA_700M, KOALA_1B, Segmind_Vega, SD_X4Upscaler, Stable_Cascade_C, Stable_Cascade_B] +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_extras/nodes_stable3d.py b/comfy_extras/nodes_stable3d.py index 4375d8f96..be2e34c28 100644 --- a/comfy_extras/nodes_stable3d.py +++ b/comfy_extras/nodes_stable3d.py @@ -29,8 +29,8 @@ class StableZero123_Conditioning: "width": ("INT", {"default": 256, "min": 16, "max": nodes.MAX_RESOLUTION, "step": 8}), "height": ("INT", {"default": 256, "min": 16, "max": nodes.MAX_RESOLUTION, "step": 8}), "batch_size": ("INT", {"default": 1, "min": 1, "max": 4096}), - "elevation": ("FLOAT", {"default": 0.0, "min": -180.0, "max": 180.0}), - "azimuth": ("FLOAT", {"default": 0.0, "min": -180.0, "max": 180.0}), + "elevation": ("FLOAT", {"default": 0.0, "min": -180.0, "max": 180.0, "step": 0.1, "round": False}), + "azimuth": ("FLOAT", {"default": 0.0, "min": -180.0, "max": 180.0, "step": 0.1, "round": False}), }} RETURN_TYPES = ("CONDITIONING", "CONDITIONING", "LATENT") RETURN_NAMES = ("positive", "negative", "latent") @@ -62,10 +62,10 @@ class StableZero123_Conditioning_Batched: "width": ("INT", {"default": 256, "min": 16, "max": nodes.MAX_RESOLUTION, "step": 8}), "height": ("INT", {"default": 256, "min": 16, "max": nodes.MAX_RESOLUTION, "step": 8}), "batch_size": ("INT", {"default": 1, "min": 1, "max": 4096}), - "elevation": ("FLOAT", {"default": 0.0, "min": -180.0, "max": 180.0}), - "azimuth": ("FLOAT", {"default": 0.0, "min": -180.0, "max": 180.0}), - "elevation_batch_increment": ("FLOAT", {"default": 0.0, "min": -180.0, "max": 180.0}), - "azimuth_batch_increment": ("FLOAT", {"default": 0.0, "min": -180.0, "max": 180.0}), + "elevation": ("FLOAT", {"default": 0.0, "min": -180.0, "max": 180.0, "step": 0.1, "round": False}), + "azimuth": ("FLOAT", {"default": 0.0, "min": -180.0, "max": 180.0, "step": 0.1, "round": False}), + "elevation_batch_increment": ("FLOAT", {"default": 0.0, "min": -180.0, "max": 180.0, "step": 0.1, "round": False}), + "azimuth_batch_increment": ("FLOAT", {"default": 0.0, "min": -180.0, "max": 180.0, "step": 0.1, "round": False}), }} RETURN_TYPES = ("CONDITIONING", "CONDITIONING", "LATENT") RETURN_NAMES = ("positive", "negative", "latent") @@ -95,8 +95,49 @@ class StableZero123_Conditioning_Batched: latent = torch.zeros([batch_size, 4, height // 8, width // 8]) return (positive, negative, {"samples":latent, "batch_index": [0] * batch_size}) +class SV3D_Conditioning: + @classmethod + def INPUT_TYPES(s): + return {"required": { "clip_vision": ("CLIP_VISION",), + "init_image": ("IMAGE",), + "vae": ("VAE",), + "width": ("INT", {"default": 576, "min": 16, "max": nodes.MAX_RESOLUTION, "step": 8}), + "height": ("INT", {"default": 576, "min": 16, "max": nodes.MAX_RESOLUTION, "step": 8}), + "video_frames": ("INT", {"default": 21, "min": 1, "max": 4096}), + "elevation": ("FLOAT", {"default": 0.0, "min": -90.0, "max": 90.0, "step": 0.1, "round": False}), + }} + RETURN_TYPES = ("CONDITIONING", "CONDITIONING", "LATENT") + RETURN_NAMES = ("positive", "negative", "latent") + + FUNCTION = "encode" + + CATEGORY = "conditioning/3d_models" + + def encode(self, clip_vision, init_image, vae, width, height, video_frames, elevation): + output = clip_vision.encode_image(init_image) + pooled = output.image_embeds.unsqueeze(0) + pixels = comfy.utils.common_upscale(init_image.movedim(-1,1), width, height, "bilinear", "center").movedim(1,-1) + encode_pixels = pixels[:,:,:,:3] + t = vae.encode(encode_pixels) + + azimuth = 0 + azimuth_increment = 360 / (max(video_frames, 2) - 1) + + elevations = [] + azimuths = [] + for i in range(video_frames): + elevations.append(elevation) + azimuths.append(azimuth) + azimuth += azimuth_increment + + positive = [[pooled, {"concat_latent_image": t, "elevation": elevations, "azimuth": azimuths}]] + negative = [[torch.zeros_like(pooled), {"concat_latent_image": torch.zeros_like(t), "elevation": elevations, "azimuth": azimuths}]] + latent = torch.zeros([video_frames, 4, height // 8, width // 8]) + return (positive, negative, {"samples":latent}) + NODE_CLASS_MAPPINGS = { "StableZero123_Conditioning": StableZero123_Conditioning, "StableZero123_Conditioning_Batched": StableZero123_Conditioning_Batched, + "SV3D_Conditioning": SV3D_Conditioning, } diff --git a/comfy_extras/nodes_video_model.py b/comfy_extras/nodes_video_model.py index a52625652..1a0189ed4 100644 --- a/comfy_extras/nodes_video_model.py +++ b/comfy_extras/nodes_video_model.py @@ -79,6 +79,33 @@ class VideoLinearCFGGuidance: m.set_model_sampler_cfg_function(linear_cfg) return (m, ) +class VideoTriangleCFGGuidance: + @classmethod + def INPUT_TYPES(s): + return {"required": { "model": ("MODEL",), + "min_cfg": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 100.0, "step":0.5, "round": 0.01}), + }} + RETURN_TYPES = ("MODEL",) + FUNCTION = "patch" + + CATEGORY = "sampling/video_models" + + def patch(self, model, min_cfg): + def linear_cfg(args): + cond = args["cond"] + uncond = args["uncond"] + cond_scale = args["cond_scale"] + period = 1.0 + values = torch.linspace(0, 1, cond.shape[0], device=cond.device) + values = 2 * (values / period - torch.floor(values / period + 0.5)).abs() + scale = (values * (cond_scale - min_cfg) + min_cfg).reshape((cond.shape[0], 1, 1, 1)) + + return uncond + scale * (cond - uncond) + + m = model.clone() + m.set_model_sampler_cfg_function(linear_cfg) + return (m, ) + class ImageOnlyCheckpointSave(comfy_extras.nodes_model_merging.CheckpointSave): CATEGORY = "_for_testing" @@ -98,6 +125,7 @@ NODE_CLASS_MAPPINGS = { "ImageOnlyCheckpointLoader": ImageOnlyCheckpointLoader, "SVD_img2vid_Conditioning": SVD_img2vid_Conditioning, "VideoLinearCFGGuidance": VideoLinearCFGGuidance, + "VideoTriangleCFGGuidance": VideoTriangleCFGGuidance, "ImageOnlyCheckpointSave": ImageOnlyCheckpointSave, } From 0c55f16c9e66eaa4915e288b34e4f848fb2d949f Mon Sep 17 00:00:00 2001 From: comfyanonymous Date: Tue, 19 Mar 2024 02:34:01 -0400 Subject: [PATCH 08/54] Remove code that should be useless now. --- nodes.py | 3 --- 1 file changed, 3 deletions(-) diff --git a/nodes.py b/nodes.py index 453f6e606..d24df200c 100644 --- a/nodes.py +++ b/nodes.py @@ -15,9 +15,6 @@ from PIL.PngImagePlugin import PngInfo import numpy as np import safetensors.torch -sys.path.insert(0, os.path.join(os.path.dirname(os.path.realpath(__file__)), "comfy")) - - import comfy.diffusers_load import comfy.samplers import comfy.sample From d14bdb18967f7413852a364747c49599de537eec Mon Sep 17 00:00:00 2001 From: comfyanonymous Date: Tue, 19 Mar 2024 11:17:37 -0400 Subject: [PATCH 09/54] Revert, NOTE: this will be removed again soon please fix your nodes. --- nodes.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/nodes.py b/nodes.py index d24df200c..453f6e606 100644 --- a/nodes.py +++ b/nodes.py @@ -15,6 +15,9 @@ from PIL.PngImagePlugin import PngInfo import numpy as np import safetensors.torch +sys.path.insert(0, os.path.join(os.path.dirname(os.path.realpath(__file__)), "comfy")) + + import comfy.diffusers_load import comfy.samplers import comfy.sample From 150a3e946fe49990454b0b81e26f06cdd9fbab8f Mon Sep 17 00:00:00 2001 From: comfyanonymous Date: Wed, 20 Mar 2024 01:35:59 -0400 Subject: [PATCH 10/54] Make LCM sampler use the model noise scaling function. --- comfy/k_diffusion/sampling.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/comfy/k_diffusion/sampling.py b/comfy/k_diffusion/sampling.py index 761c2e0ef..57518c7b6 100644 --- a/comfy/k_diffusion/sampling.py +++ b/comfy/k_diffusion/sampling.py @@ -748,7 +748,7 @@ def sample_lcm(model, x, sigmas, extra_args=None, callback=None, disable=None, n x = denoised if sigmas[i + 1] > 0: - x += sigmas[i + 1] * noise_sampler(sigmas[i], sigmas[i + 1]) + x = model.inner_model.inner_model.model_sampling.noise_scaling(sigmas[i + 1] * s_in, noise_sampler(sigmas[i], sigmas[i + 1]), x) return x From c18a203a8abdd0fce24743a838fe0d0400d8ff09 Mon Sep 17 00:00:00 2001 From: comfyanonymous Date: Wed, 20 Mar 2024 01:29:26 -0400 Subject: [PATCH 11/54] Don't unload model weights for non weight patches. --- comfy/model_management.py | 44 ++++++++++++++++++++++------ comfy/model_patcher.py | 60 ++++++++++++++++++++++++++------------- 2 files changed, 76 insertions(+), 28 deletions(-) diff --git a/comfy/model_management.py b/comfy/model_management.py index 66fa918bb..74958908a 100644 --- a/comfy/model_management.py +++ b/comfy/model_management.py @@ -273,6 +273,7 @@ class LoadedModel: def __init__(self, model): self.model = model self.device = model.load_device + self.weights_loaded = False def model_memory(self): return self.model.model_size() @@ -289,11 +290,13 @@ class LoadedModel: self.model.model_patches_to(self.device) self.model.model_patches_to(self.model.model_dtype()) + load_weights = not self.weights_loaded + try: - if lowvram_model_memory > 0: + if lowvram_model_memory > 0 and load_weights: self.real_model = self.model.patch_model_lowvram(device_to=patch_model_to, lowvram_model_memory=lowvram_model_memory) else: - self.real_model = self.model.patch_model(device_to=patch_model_to) + self.real_model = self.model.patch_model(device_to=patch_model_to, patch_weights=load_weights) except Exception as e: self.model.unpatch_model(self.model.offload_device) self.model_unload() @@ -302,11 +305,13 @@ class LoadedModel: if is_intel_xpu() and not args.disable_ipex_optimize: self.real_model = torch.xpu.optimize(self.real_model.eval(), inplace=True, auto_kernel_selection=True, graph_mode=True) + self.weights_loaded = True return self.real_model - def model_unload(self): - self.model.unpatch_model(self.model.offload_device) + def model_unload(self, unpatch_weights=True): + self.model.unpatch_model(self.model.offload_device, unpatch_weights=unpatch_weights) self.model.model_patches_to(self.model.offload_device) + self.weights_loaded = self.weights_loaded and not unpatch_weights def __eq__(self, other): return self.model is other.model @@ -314,15 +319,35 @@ class LoadedModel: def minimum_inference_memory(): return (1024 * 1024 * 1024) -def unload_model_clones(model): +def unload_model_clones(loaded_model, unload_weights_only=True): + model = loaded_model.model + to_unload = [] for i in range(len(current_loaded_models)): if model.is_clone(current_loaded_models[i].model): to_unload = [i] + to_unload + if len(to_unload) == 0: + return + + same_weights = 0 for i in to_unload: - logging.debug("unload clone {}".format(i)) - current_loaded_models.pop(i).model_unload() + if model.clone_has_same_weights(current_loaded_models[i].model): + same_weights += 1 + + if same_weights == len(to_unload): + unload_weight = False + else: + unload_weight = True + + if unload_weights_only and unload_weight == False: + return + + for i in to_unload: + logging.debug("unload clone {} {}".format(i, unload_weight)) + current_loaded_models.pop(i).model_unload(unpatch_weights=unload_weight) + + loaded_model.weights_loaded = not unload_weight def free_memory(memory_required, device, keep_loaded=[]): unloaded_model = False @@ -377,13 +402,16 @@ def load_models_gpu(models, memory_required=0): total_memory_required = {} for loaded_model in models_to_load: - unload_model_clones(loaded_model.model) + unload_model_clones(loaded_model, unload_weights_only=True) #unload clones where the weights are different total_memory_required[loaded_model.device] = total_memory_required.get(loaded_model.device, 0) + loaded_model.model_memory_required(loaded_model.device) for device in total_memory_required: if device != torch.device("cpu"): free_memory(total_memory_required[device] * 1.3 + extra_mem, device, models_already_loaded) + for loaded_model in models_to_load: + unload_model_clones(loaded_model, unload_weights_only=False) #unload the rest of the clones where the weights can stay loaded + for loaded_model in models_to_load: model = loaded_model.model torch_dev = model.load_device diff --git a/comfy/model_patcher.py b/comfy/model_patcher.py index 475fa812c..aa78302d2 100644 --- a/comfy/model_patcher.py +++ b/comfy/model_patcher.py @@ -2,6 +2,7 @@ import torch import copy import inspect import logging +import uuid import comfy.utils import comfy.model_management @@ -25,6 +26,7 @@ class ModelPatcher: self.weight_inplace_update = weight_inplace_update self.model_lowvram = False + self.patches_uuid = uuid.uuid4() def model_size(self): if self.size > 0: @@ -39,10 +41,13 @@ class ModelPatcher: n.patches = {} for k in self.patches: n.patches[k] = self.patches[k][:] + n.patches_uuid = self.patches_uuid n.object_patches = self.object_patches.copy() n.model_options = copy.deepcopy(self.model_options) n.model_keys = self.model_keys + n.backup = self.backup + n.object_patches_backup = self.object_patches_backup return n def is_clone(self, other): @@ -50,6 +55,19 @@ class ModelPatcher: return True return False + def clone_has_same_weights(self, clone): + if not self.is_clone(clone): + return False + + if len(self.patches) == 0 and len(clone.patches) == 0: + return True + + if self.patches_uuid == clone.patches_uuid: + if len(self.patches) != len(clone.patches): + logging.warning("WARNING: something went wrong, same patch uuid but different length of patches.") + else: + return True + def memory_required(self, input_shape): return self.model.memory_required(input_shape=input_shape) @@ -154,6 +172,7 @@ class ModelPatcher: current_patches.append((strength_patch, patches[k], strength_model)) self.patches[k] = current_patches + self.patches_uuid = uuid.uuid4() return list(p) def get_key_patches(self, filter_prefix=None): @@ -387,31 +406,32 @@ class ModelPatcher: return weight - def unpatch_model(self, device_to=None): - if self.model_lowvram: - for m in self.model.modules(): - if hasattr(m, "prev_comfy_cast_weights"): - m.comfy_cast_weights = m.prev_comfy_cast_weights - del m.prev_comfy_cast_weights - m.weight_function = None - m.bias_function = None + def unpatch_model(self, device_to=None, unpatch_weights=True): + if unpatch_weights: + if self.model_lowvram: + for m in self.model.modules(): + if hasattr(m, "prev_comfy_cast_weights"): + m.comfy_cast_weights = m.prev_comfy_cast_weights + del m.prev_comfy_cast_weights + m.weight_function = None + m.bias_function = None - self.model_lowvram = False + self.model_lowvram = False - keys = list(self.backup.keys()) + keys = list(self.backup.keys()) - if self.weight_inplace_update: - for k in keys: - comfy.utils.copy_to_param(self.model, k, self.backup[k]) - else: - for k in keys: - comfy.utils.set_attr_param(self.model, k, self.backup[k]) + if self.weight_inplace_update: + for k in keys: + comfy.utils.copy_to_param(self.model, k, self.backup[k]) + else: + for k in keys: + comfy.utils.set_attr_param(self.model, k, self.backup[k]) - self.backup = {} + self.backup.clear() - if device_to is not None: - self.model.to(device_to) - self.current_device = device_to + if device_to is not None: + self.model.to(device_to) + self.current_device = device_to keys = list(self.object_patches_backup.keys()) for k in keys: From 4b9005e949224782236a8b914eae48bc503f1f18 Mon Sep 17 00:00:00 2001 From: comfyanonymous Date: Wed, 20 Mar 2024 13:53:45 -0400 Subject: [PATCH 12/54] Fix regression with model merging. --- comfy/model_management.py | 19 ++++++++++--------- 1 file changed, 10 insertions(+), 9 deletions(-) diff --git a/comfy/model_management.py b/comfy/model_management.py index 74958908a..11c97f290 100644 --- a/comfy/model_management.py +++ b/comfy/model_management.py @@ -319,16 +319,14 @@ class LoadedModel: def minimum_inference_memory(): return (1024 * 1024 * 1024) -def unload_model_clones(loaded_model, unload_weights_only=True): - model = loaded_model.model - +def unload_model_clones(model, unload_weights_only=True, force_unload=True): to_unload = [] for i in range(len(current_loaded_models)): if model.is_clone(current_loaded_models[i].model): to_unload = [i] + to_unload if len(to_unload) == 0: - return + return None same_weights = 0 for i in to_unload: @@ -340,14 +338,15 @@ def unload_model_clones(loaded_model, unload_weights_only=True): else: unload_weight = True - if unload_weights_only and unload_weight == False: - return + if not force_unload: + if unload_weights_only and unload_weight == False: + return None for i in to_unload: logging.debug("unload clone {} {}".format(i, unload_weight)) current_loaded_models.pop(i).model_unload(unpatch_weights=unload_weight) - loaded_model.weights_loaded = not unload_weight + return unload_weight def free_memory(memory_required, device, keep_loaded=[]): unloaded_model = False @@ -402,7 +401,7 @@ def load_models_gpu(models, memory_required=0): total_memory_required = {} for loaded_model in models_to_load: - unload_model_clones(loaded_model, unload_weights_only=True) #unload clones where the weights are different + unload_model_clones(loaded_model.model, unload_weights_only=True, force_unload=False) #unload clones where the weights are different total_memory_required[loaded_model.device] = total_memory_required.get(loaded_model.device, 0) + loaded_model.model_memory_required(loaded_model.device) for device in total_memory_required: @@ -410,7 +409,9 @@ def load_models_gpu(models, memory_required=0): free_memory(total_memory_required[device] * 1.3 + extra_mem, device, models_already_loaded) for loaded_model in models_to_load: - unload_model_clones(loaded_model, unload_weights_only=False) #unload the rest of the clones where the weights can stay loaded + weights_unloaded = unload_model_clones(loaded_model.model, unload_weights_only=False, force_unload=False) #unload the rest of the clones where the weights can stay loaded + if weights_unloaded is not None: + loaded_model.weights_loaded = not weights_unloaded for loaded_model in models_to_load: model = loaded_model.model From 5d875d77fe6e31a4b0bc6dc36f0441eba3b6afe1 Mon Sep 17 00:00:00 2001 From: comfyanonymous Date: Wed, 20 Mar 2024 20:48:54 -0400 Subject: [PATCH 13/54] Fix regression with lcm not working with batches. --- comfy/k_diffusion/sampling.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/comfy/k_diffusion/sampling.py b/comfy/k_diffusion/sampling.py index 57518c7b6..7af016829 100644 --- a/comfy/k_diffusion/sampling.py +++ b/comfy/k_diffusion/sampling.py @@ -748,7 +748,7 @@ def sample_lcm(model, x, sigmas, extra_args=None, callback=None, disable=None, n x = denoised if sigmas[i + 1] > 0: - x = model.inner_model.inner_model.model_sampling.noise_scaling(sigmas[i + 1] * s_in, noise_sampler(sigmas[i], sigmas[i + 1]), x) + x = model.inner_model.inner_model.model_sampling.noise_scaling(sigmas[i + 1], noise_sampler(sigmas[i], sigmas[i + 1]), x) return x From 062483823738ed610d8d074ba63910c90e9d45b7 Mon Sep 17 00:00:00 2001 From: comfyanonymous Date: Thu, 21 Mar 2024 14:49:11 -0400 Subject: [PATCH 14/54] Add inverse noise scaling function. --- comfy/model_sampling.py | 3 +++ comfy/samplers.py | 1 + 2 files changed, 4 insertions(+) diff --git a/comfy/model_sampling.py b/comfy/model_sampling.py index d325f76d9..37976b326 100644 --- a/comfy/model_sampling.py +++ b/comfy/model_sampling.py @@ -20,6 +20,9 @@ class EPS: noise += latent_image return noise + def inverse_noise_scaling(self, sigma, latent): + return latent + class V_PREDICTION(EPS): def calculate_denoised(self, sigma, model_output, model_input): sigma = sigma.view(sigma.shape[:1] + (1,) * (model_output.ndim - 1)) diff --git a/comfy/samplers.py b/comfy/samplers.py index d721cb2e5..3678dc818 100644 --- a/comfy/samplers.py +++ b/comfy/samplers.py @@ -546,6 +546,7 @@ class KSAMPLER(Sampler): k_callback = lambda x: callback(x["i"], x["denoised"], x["x"], total_steps) samples = self.sampler_function(model_k, noise, sigmas, extra_args=extra_args, callback=k_callback, disable=disable_pbar, **self.extra_options) + samples = model_wrap.inner_model.model_sampling.inverse_noise_scaling(sigmas[-1], samples) return samples From a28a9dc83684624ee2167c0b92d976bb68f2c606 Mon Sep 17 00:00:00 2001 From: comfyanonymous Date: Fri, 22 Mar 2024 12:56:48 -0400 Subject: [PATCH 15/54] Add an example to use the SaveImageWebsocket node and enable it. --- ...ve.py.disabled => websocket_image_save.py} | 6 +- .../websockets_api_example_ws_images.py | 159 ++++++++++++++++++ 2 files changed, 160 insertions(+), 5 deletions(-) rename custom_nodes/{websocket_image_save.py.disabled => websocket_image_save.py} (84%) create mode 100644 script_examples/websockets_api_example_ws_images.py diff --git a/custom_nodes/websocket_image_save.py.disabled b/custom_nodes/websocket_image_save.py similarity index 84% rename from custom_nodes/websocket_image_save.py.disabled rename to custom_nodes/websocket_image_save.py index b85a5de8b..5aa573642 100644 --- a/custom_nodes/websocket_image_save.py.disabled +++ b/custom_nodes/websocket_image_save.py @@ -10,10 +10,6 @@ import time #binary images on the websocket with a 8 byte header indicating the type #of binary message (first 4 bytes) and the image format (next 4 bytes). -#The reason this node is disabled by default is because there is a small -#issue when using it with the default ComfyUI web interface: When generating -#batches only the last image will be shown in the UI. - #Note that no metadata will be put in the images saved with this node. class SaveImageWebsocket: @@ -28,7 +24,7 @@ class SaveImageWebsocket: OUTPUT_NODE = True - CATEGORY = "image" + CATEGORY = "api/image" def save_images(self, images): pbar = comfy.utils.ProgressBar(images.shape[0]) diff --git a/script_examples/websockets_api_example_ws_images.py b/script_examples/websockets_api_example_ws_images.py new file mode 100644 index 000000000..737488621 --- /dev/null +++ b/script_examples/websockets_api_example_ws_images.py @@ -0,0 +1,159 @@ +#This is an example that uses the websockets api and the SaveImageWebsocket node to get images directly without +#them being saved to disk + +import websocket #NOTE: websocket-client (https://github.com/websocket-client/websocket-client) +import uuid +import json +import urllib.request +import urllib.parse + +server_address = "127.0.0.1:8188" +client_id = str(uuid.uuid4()) + +def queue_prompt(prompt): + p = {"prompt": prompt, "client_id": client_id} + data = json.dumps(p).encode('utf-8') + req = urllib.request.Request("http://{}/prompt".format(server_address), data=data) + return json.loads(urllib.request.urlopen(req).read()) + +def get_image(filename, subfolder, folder_type): + data = {"filename": filename, "subfolder": subfolder, "type": folder_type} + url_values = urllib.parse.urlencode(data) + with urllib.request.urlopen("http://{}/view?{}".format(server_address, url_values)) as response: + return response.read() + +def get_history(prompt_id): + with urllib.request.urlopen("http://{}/history/{}".format(server_address, prompt_id)) as response: + return json.loads(response.read()) + +def get_images(ws, prompt): + prompt_id = queue_prompt(prompt)['prompt_id'] + output_images = {} + current_node = "" + while True: + out = ws.recv() + if isinstance(out, str): + message = json.loads(out) + if message['type'] == 'executing': + data = message['data'] + if data['prompt_id'] == prompt_id: + if data['node'] is None: + break #Execution is done + else: + current_node = data['node'] + else: + if current_node == 'save_image_websocket_node': + images_output = output_images.get(current_node, []) + images_output.append(out[8:]) + output_images[current_node] = images_output + + return output_images + +prompt_text = """ +{ + "3": { + "class_type": "KSampler", + "inputs": { + "cfg": 8, + "denoise": 1, + "latent_image": [ + "5", + 0 + ], + "model": [ + "4", + 0 + ], + "negative": [ + "7", + 0 + ], + "positive": [ + "6", + 0 + ], + "sampler_name": "euler", + "scheduler": "normal", + "seed": 8566257, + "steps": 20 + } + }, + "4": { + "class_type": "CheckpointLoaderSimple", + "inputs": { + "ckpt_name": "v1-5-pruned-emaonly.ckpt" + } + }, + "5": { + "class_type": "EmptyLatentImage", + "inputs": { + "batch_size": 1, + "height": 512, + "width": 512 + } + }, + "6": { + "class_type": "CLIPTextEncode", + "inputs": { + "clip": [ + "4", + 1 + ], + "text": "masterpiece best quality girl" + } + }, + "7": { + "class_type": "CLIPTextEncode", + "inputs": { + "clip": [ + "4", + 1 + ], + "text": "bad hands" + } + }, + "8": { + "class_type": "VAEDecode", + "inputs": { + "samples": [ + "3", + 0 + ], + "vae": [ + "4", + 2 + ] + } + }, + "save_image_websocket_node": { + "class_type": "SaveImageWebsocket", + "inputs": { + "images": [ + "8", + 0 + ] + } + } +} +""" + +prompt = json.loads(prompt_text) +#set the text prompt for our positive CLIPTextEncode +prompt["6"]["inputs"]["text"] = "masterpiece best quality man" + +#set the seed for our KSampler node +prompt["3"]["inputs"]["seed"] = 5 + +ws = websocket.WebSocket() +ws.connect("ws://{}/ws?clientId={}".format(server_address, client_id)) +images = get_images(ws, prompt) + +#Commented out code to display the output images: + +# for node_id in images: +# for image_data in images[node_id]: +# from PIL import Image +# import io +# image = Image.open(io.BytesIO(image_data)) +# image.show() + From 6a32c06f0646805456ac2744ece289962a2379c6 Mon Sep 17 00:00:00 2001 From: comfyanonymous Date: Sat, 23 Mar 2024 17:27:10 -0400 Subject: [PATCH 16/54] Move cleanup_models to improve performance. --- execution.py | 1 - main.py | 1 + 2 files changed, 1 insertion(+), 1 deletion(-) diff --git a/execution.py b/execution.py index 5f8ea835c..1b8f606a5 100644 --- a/execution.py +++ b/execution.py @@ -368,7 +368,6 @@ class PromptExecutor: d = self.outputs_ui.pop(x) del d - comfy.model_management.cleanup_models() self.add_message("execution_cached", { "nodes": list(current_outputs) , "prompt_id": prompt_id}, broadcast=False) diff --git a/main.py b/main.py index 3dee72e3a..b3a3ebea8 100644 --- a/main.py +++ b/main.py @@ -139,6 +139,7 @@ def prompt_worker(q, server): if need_gc: current_time = time.perf_counter() if (current_time - last_gc_collect) > gc_collect_interval: + comfy.model_management.cleanup_models() gc.collect() comfy.model_management.soft_empty_cache() last_gc_collect = current_time From c6de09b02e20d748739fb6af58e196ebdd12825a Mon Sep 17 00:00:00 2001 From: comfyanonymous Date: Sun, 24 Mar 2024 02:36:30 -0400 Subject: [PATCH 17/54] Optimize memory unload strategy for more optimized performance. --- comfy/model_management.py | 25 ++++++++++++++++--------- 1 file changed, 16 insertions(+), 9 deletions(-) diff --git a/comfy/model_management.py b/comfy/model_management.py index 11c97f290..715ca2ee2 100644 --- a/comfy/model_management.py +++ b/comfy/model_management.py @@ -349,20 +349,27 @@ def unload_model_clones(model, unload_weights_only=True, force_unload=True): return unload_weight def free_memory(memory_required, device, keep_loaded=[]): - unloaded_model = False + unloaded_model = [] + can_unload = [] + for i in range(len(current_loaded_models) -1, -1, -1): - if not DISABLE_SMART_MEMORY: - if get_free_memory(device) > memory_required: - break shift_model = current_loaded_models[i] if shift_model.device == device: if shift_model not in keep_loaded: - m = current_loaded_models.pop(i) - m.model_unload() - del m - unloaded_model = True + can_unload.append((sys.getrefcount(shift_model.model), shift_model.model_memory(), i)) - if unloaded_model: + for x in sorted(can_unload): + i = x[-1] + if not DISABLE_SMART_MEMORY: + if get_free_memory(device) > memory_required: + break + current_loaded_models[i].model_unload() + unloaded_model.append(i) + + for i in sorted(unloaded_model, reverse=True): + current_loaded_models.pop(i) + + if len(unloaded_model) > 0: soft_empty_cache() else: if vram_state != VRAMState.HIGH_VRAM: From ae77590b4e0cbab55ac9edb67722092d375fe3ec Mon Sep 17 00:00:00 2001 From: comfyanonymous Date: Mon, 25 Mar 2024 18:09:23 -0400 Subject: [PATCH 18/54] dora_scale support for lora file. --- comfy/lora.py | 14 ++++++++++---- comfy/model_patcher.py | 25 +++++++++++++++++++++++++ 2 files changed, 35 insertions(+), 4 deletions(-) diff --git a/comfy/lora.py b/comfy/lora.py index 637380d54..096285bba 100644 --- a/comfy/lora.py +++ b/comfy/lora.py @@ -21,6 +21,12 @@ def load_lora(lora, to_load): alpha = lora[alpha_name].item() loaded_keys.add(alpha_name) + dora_scale_name = "{}.dora_scale".format(x) + dora_scale = None + if dora_scale_name in lora.keys(): + dora_scale = lora[dora_scale_name] + loaded_keys.add(dora_scale_name) + regular_lora = "{}.lora_up.weight".format(x) diffusers_lora = "{}_lora.up.weight".format(x) transformers_lora = "{}.lora_linear_layer.up.weight".format(x) @@ -44,7 +50,7 @@ def load_lora(lora, to_load): if mid_name is not None and mid_name in lora.keys(): mid = lora[mid_name] loaded_keys.add(mid_name) - patch_dict[to_load[x]] = ("lora", (lora[A_name], lora[B_name], alpha, mid)) + patch_dict[to_load[x]] = ("lora", (lora[A_name], lora[B_name], alpha, mid, dora_scale)) loaded_keys.add(A_name) loaded_keys.add(B_name) @@ -65,7 +71,7 @@ def load_lora(lora, to_load): loaded_keys.add(hada_t1_name) loaded_keys.add(hada_t2_name) - patch_dict[to_load[x]] = ("loha", (lora[hada_w1_a_name], lora[hada_w1_b_name], alpha, lora[hada_w2_a_name], lora[hada_w2_b_name], hada_t1, hada_t2)) + patch_dict[to_load[x]] = ("loha", (lora[hada_w1_a_name], lora[hada_w1_b_name], alpha, lora[hada_w2_a_name], lora[hada_w2_b_name], hada_t1, hada_t2, dora_scale)) loaded_keys.add(hada_w1_a_name) loaded_keys.add(hada_w1_b_name) loaded_keys.add(hada_w2_a_name) @@ -117,7 +123,7 @@ def load_lora(lora, to_load): loaded_keys.add(lokr_t2_name) if (lokr_w1 is not None) or (lokr_w2 is not None) or (lokr_w1_a is not None) or (lokr_w2_a is not None): - patch_dict[to_load[x]] = ("lokr", (lokr_w1, lokr_w2, alpha, lokr_w1_a, lokr_w1_b, lokr_w2_a, lokr_w2_b, lokr_t2)) + patch_dict[to_load[x]] = ("lokr", (lokr_w1, lokr_w2, alpha, lokr_w1_a, lokr_w1_b, lokr_w2_a, lokr_w2_b, lokr_t2, dora_scale)) #glora a1_name = "{}.a1.weight".format(x) @@ -125,7 +131,7 @@ def load_lora(lora, to_load): b1_name = "{}.b1.weight".format(x) b2_name = "{}.b2.weight".format(x) if a1_name in lora: - patch_dict[to_load[x]] = ("glora", (lora[a1_name], lora[a2_name], lora[b1_name], lora[b2_name], alpha)) + patch_dict[to_load[x]] = ("glora", (lora[a1_name], lora[a2_name], lora[b1_name], lora[b2_name], alpha, dora_scale)) loaded_keys.add(a1_name) loaded_keys.add(a2_name) loaded_keys.add(b1_name) diff --git a/comfy/model_patcher.py b/comfy/model_patcher.py index aa78302d2..8dda84cfd 100644 --- a/comfy/model_patcher.py +++ b/comfy/model_patcher.py @@ -7,6 +7,18 @@ import uuid import comfy.utils import comfy.model_management +def apply_weight_decompose(dora_scale, weight): + weight_norm = ( + weight.transpose(0, 1) + .reshape(weight.shape[1], -1) + .norm(dim=1, keepdim=True) + .reshape(weight.shape[1], *[1] * (weight.dim() - 1)) + .transpose(0, 1) + ) + + return weight * (dora_scale / weight_norm) + + class ModelPatcher: def __init__(self, model, load_device, offload_device, size=0, current_device=None, weight_inplace_update=False): self.size = size @@ -309,6 +321,7 @@ class ModelPatcher: elif patch_type == "lora": #lora/locon mat1 = comfy.model_management.cast_to_device(v[0], weight.device, torch.float32) mat2 = comfy.model_management.cast_to_device(v[1], weight.device, torch.float32) + dora_scale = v[4] if v[2] is not None: alpha *= v[2] / mat2.shape[0] if v[3] is not None: @@ -318,6 +331,8 @@ class ModelPatcher: mat2 = torch.mm(mat2.transpose(0, 1).flatten(start_dim=1), mat3.transpose(0, 1).flatten(start_dim=1)).reshape(final_shape).transpose(0, 1) try: weight += (alpha * torch.mm(mat1.flatten(start_dim=1), mat2.flatten(start_dim=1))).reshape(weight.shape).type(weight.dtype) + if dora_scale is not None: + weight = apply_weight_decompose(comfy.model_management.cast_to_device(dora_scale, weight.device, torch.float32), weight) except Exception as e: logging.error("ERROR {} {} {}".format(patch_type, key, e)) elif patch_type == "lokr": @@ -328,6 +343,7 @@ class ModelPatcher: w2_a = v[5] w2_b = v[6] t2 = v[7] + dora_scale = v[8] dim = None if w1 is None: @@ -357,6 +373,8 @@ class ModelPatcher: try: weight += alpha * torch.kron(w1, w2).reshape(weight.shape).type(weight.dtype) + if dora_scale is not None: + weight = apply_weight_decompose(comfy.model_management.cast_to_device(dora_scale, weight.device, torch.float32), weight) except Exception as e: logging.error("ERROR {} {} {}".format(patch_type, key, e)) elif patch_type == "loha": @@ -366,6 +384,7 @@ class ModelPatcher: alpha *= v[2] / w1b.shape[0] w2a = v[3] w2b = v[4] + dora_scale = v[7] if v[5] is not None: #cp decomposition t1 = v[5] t2 = v[6] @@ -386,12 +405,16 @@ class ModelPatcher: try: weight += (alpha * m1 * m2).reshape(weight.shape).type(weight.dtype) + if dora_scale is not None: + weight = apply_weight_decompose(comfy.model_management.cast_to_device(dora_scale, weight.device, torch.float32), weight) except Exception as e: logging.error("ERROR {} {} {}".format(patch_type, key, e)) elif patch_type == "glora": if v[4] is not None: alpha *= v[4] / v[0].shape[0] + dora_scale = v[5] + a1 = comfy.model_management.cast_to_device(v[0].flatten(start_dim=1), weight.device, torch.float32) a2 = comfy.model_management.cast_to_device(v[1].flatten(start_dim=1), weight.device, torch.float32) b1 = comfy.model_management.cast_to_device(v[2].flatten(start_dim=1), weight.device, torch.float32) @@ -399,6 +422,8 @@ class ModelPatcher: try: weight += ((torch.mm(b2, b1) + torch.mm(torch.mm(weight.flatten(start_dim=1), a2), a1)) * alpha).reshape(weight.shape).type(weight.dtype) + if dora_scale is not None: + weight = apply_weight_decompose(comfy.model_management.cast_to_device(dora_scale, weight.device, torch.float32), weight) except Exception as e: logging.error("ERROR {} {} {}".format(patch_type, key, e)) else: From 11838e60f4aac495c6f0aed09415f3e8e0d2a402 Mon Sep 17 00:00:00 2001 From: comfyanonymous Date: Tue, 26 Mar 2024 04:00:53 -0400 Subject: [PATCH 19/54] Increase the max resolution. --- nodes.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/nodes.py b/nodes.py index 453f6e606..6c05e3a2b 100644 --- a/nodes.py +++ b/nodes.py @@ -41,7 +41,7 @@ def before_node_execution(): def interrupt_processing(value=True): comfy.model_management.interrupt_current_processing(value) -MAX_RESOLUTION=8192 +MAX_RESOLUTION=16384 class CLIPTextEncode: @classmethod From c9673926fb2df5e82ccbfe2a746c428c148450f0 Mon Sep 17 00:00:00 2001 From: comfyanonymous Date: Tue, 26 Mar 2024 04:07:30 -0400 Subject: [PATCH 20/54] Fix test. --- tests-ui/tests/groupNode.test.js | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests-ui/tests/groupNode.test.js b/tests-ui/tests/groupNode.test.js index e6ebedd91..53a5828d1 100644 --- a/tests-ui/tests/groupNode.test.js +++ b/tests-ui/tests/groupNode.test.js @@ -947,7 +947,7 @@ describe("group node", () => { expect(p1.widgets.value.widget.options?.step).toBe(80); // width/height step * 10 expect(p2.widgets.value.widget.options?.min).toBe(16); // width/height min - expect(p2.widgets.value.widget.options?.max).toBe(8192); // width/height max + expect(p2.widgets.value.widget.options?.max).toBe(16384); // width/height max expect(p2.widgets.value.widget.options?.step).toBe(80); // width/height step * 10 expect(p1.widgets.value.value).toBe(128); From 2f93b9164636329b9e02fe8780cdd96c1300b90d Mon Sep 17 00:00:00 2001 From: comfyanonymous Date: Tue, 26 Mar 2024 19:33:40 -0400 Subject: [PATCH 21/54] Add Tesla GPUs to cuda malloc blacklist. --- cuda_malloc.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/cuda_malloc.py b/cuda_malloc.py index 70e7ecf9a..eb2857c5f 100644 --- a/cuda_malloc.py +++ b/cuda_malloc.py @@ -47,7 +47,7 @@ blacklist = {"GeForce GTX TITAN X", "GeForce GTX 980", "GeForce GTX 970", "GeFor "Quadro K1200", "Quadro K2200", "Quadro M500", "Quadro M520", "Quadro M600", "Quadro M620", "Quadro M1000", "Quadro M1200", "Quadro M2000", "Quadro M2200", "Quadro M3000", "Quadro M4000", "Quadro M5000", "Quadro M5500", "Quadro M6000", "GeForce MX110", "GeForce MX130", "GeForce 830M", "GeForce 840M", "GeForce GTX 850M", "GeForce GTX 860M", - "GeForce GTX 1650", "GeForce GTX 1630" + "GeForce GTX 1650", "GeForce GTX 1630", "Tesla M4", "Tesla M6", "Tesla M10", "Tesla M40", "Tesla M60" } def cuda_malloc_supported(): From 8ae1e4d125100029d58f91e1938a1803abf6866c Mon Sep 17 00:00:00 2001 From: comfyanonymous Date: Wed, 27 Mar 2024 01:28:31 -0400 Subject: [PATCH 22/54] Make step on sharpen node smaller. --- comfy_extras/nodes_post_processing.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/comfy_extras/nodes_post_processing.py b/comfy_extras/nodes_post_processing.py index cb5c7d228..0110b472f 100644 --- a/comfy_extras/nodes_post_processing.py +++ b/comfy_extras/nodes_post_processing.py @@ -204,13 +204,13 @@ class Sharpen: "default": 1.0, "min": 0.1, "max": 10.0, - "step": 0.1 + "step": 0.01 }), "alpha": ("FLOAT", { "default": 1.0, "min": 0.0, "max": 5.0, - "step": 0.1 + "step": 0.01 }), }, } From 327ca1313d756c4b443790a53ab0afa1945d3f3e Mon Sep 17 00:00:00 2001 From: comfyanonymous Date: Wed, 27 Mar 2024 23:51:17 -0400 Subject: [PATCH 23/54] Support SDXS 0.9 --- comfy/model_detection.py | 8 +++++++- comfy/supported_models.py | 4 ++-- 2 files changed, 9 insertions(+), 3 deletions(-) diff --git a/comfy/model_detection.py b/comfy/model_detection.py index b7c3be309..bddbe2a49 100644 --- a/comfy/model_detection.py +++ b/comfy/model_detection.py @@ -345,7 +345,13 @@ def unet_config_from_diffusers_unet(state_dict, dtype=None): 'channel_mult': [1, 2, 4], 'transformer_depth_middle': 6, 'use_linear_in_transformer': True, 'context_dim': 2048, 'num_head_channels': 64, 'use_temporal_attention': False, 'use_temporal_resblock': 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 = {'use_checkpoint': False, 'image_size': 32, 'out_channels': 4, 'use_spatial_transformer': True, 'legacy': False, + 'adm_in_channels': None, 'dtype': dtype, 'in_channels': 4, 'model_channels': 320, 'num_res_blocks': [1, 1, 1], + 'transformer_depth': [1, 1, 1], 'channel_mult': [1, 2, 4], 'transformer_depth_middle': -2, 'use_linear_in_transformer': True, + '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] for unet_config in supported_models: matches = True diff --git a/comfy/supported_models.py b/comfy/supported_models.py index 2ce9736b7..5b2eb73fd 100644 --- a/comfy/supported_models.py +++ b/comfy/supported_models.py @@ -70,8 +70,8 @@ class SD20(supported_models_base.BASE): def model_type(self, state_dict, prefix=""): if self.unet_config["in_channels"] == 4: #SD2.0 inpainting models are not v prediction k = "{}output_blocks.11.1.transformer_blocks.0.norm1.bias".format(prefix) - out = state_dict[k] - if torch.std(out, unbiased=False) > 0.09: # not sure how well this will actually work. I guess we will find out. + out = state_dict.get(k, None) + if out is not None and torch.std(out, unbiased=False) > 0.09: # not sure how well this will actually work. I guess we will find out. return model_base.ModelType.V_PREDICTION return model_base.ModelType.EPS From 5d8898c05668b5504f8ad5bc79779381d0af35b5 Mon Sep 17 00:00:00 2001 From: comfyanonymous Date: Thu, 28 Mar 2024 18:01:04 -0400 Subject: [PATCH 24/54] Fix some performance issues with weight loading and unloading. Lower peak memory usage when changing model. Fix case where model weights would be unloaded and reloaded. --- comfy/model_management.py | 16 +++++++++++----- execution.py | 1 + 2 files changed, 12 insertions(+), 5 deletions(-) diff --git a/comfy/model_management.py b/comfy/model_management.py index 715ca2ee2..26216432a 100644 --- a/comfy/model_management.py +++ b/comfy/model_management.py @@ -274,6 +274,7 @@ class LoadedModel: self.model = model self.device = model.load_device self.weights_loaded = False + self.real_model = None def model_memory(self): return self.model.model_size() @@ -312,6 +313,7 @@ class LoadedModel: self.model.unpatch_model(self.model.offload_device, unpatch_weights=unpatch_weights) self.model.model_patches_to(self.model.offload_device) self.weights_loaded = self.weights_loaded and not unpatch_weights + self.real_model = None def __eq__(self, other): return self.model is other.model @@ -326,7 +328,7 @@ def unload_model_clones(model, unload_weights_only=True, force_unload=True): to_unload = [i] + to_unload if len(to_unload) == 0: - return None + return True same_weights = 0 for i in to_unload: @@ -408,8 +410,8 @@ def load_models_gpu(models, memory_required=0): total_memory_required = {} for loaded_model in models_to_load: - unload_model_clones(loaded_model.model, unload_weights_only=True, force_unload=False) #unload clones where the weights are different - total_memory_required[loaded_model.device] = total_memory_required.get(loaded_model.device, 0) + loaded_model.model_memory_required(loaded_model.device) + if unload_model_clones(loaded_model.model, unload_weights_only=True, force_unload=False) == True:#unload clones where the weights are different + total_memory_required[loaded_model.device] = total_memory_required.get(loaded_model.device, 0) + loaded_model.model_memory_required(loaded_model.device) for device in total_memory_required: if device != torch.device("cpu"): @@ -448,11 +450,15 @@ def load_models_gpu(models, memory_required=0): def load_model_gpu(model): return load_models_gpu([model]) -def cleanup_models(): +def cleanup_models(keep_clone_weights_loaded=False): to_delete = [] for i in range(len(current_loaded_models)): if sys.getrefcount(current_loaded_models[i].model) <= 2: - to_delete = [i] + to_delete + if not keep_clone_weights_loaded: + to_delete = [i] + to_delete + #TODO: find a less fragile way to do this. + elif sys.getrefcount(current_loaded_models[i].real_model) <= 3: #references from .real_model + the .model + to_delete = [i] + to_delete for i in to_delete: x = current_loaded_models.pop(i) diff --git a/execution.py b/execution.py index 1b8f606a5..35bdb77a6 100644 --- a/execution.py +++ b/execution.py @@ -368,6 +368,7 @@ class PromptExecutor: d = self.outputs_ui.pop(x) del d + comfy.model_management.cleanup_models(keep_clone_weights_loaded=True) self.add_message("execution_cached", { "nodes": list(current_outputs) , "prompt_id": prompt_id}, broadcast=False) From 9bf6061dfcd1370ceae3ceac07765ff84686bd20 Mon Sep 17 00:00:00 2001 From: comfyanonymous Date: Fri, 29 Mar 2024 03:07:13 -0400 Subject: [PATCH 25/54] Switch prints to logging in folder_paths and add some extra debug. --- folder_paths.py | 17 ++++++++++------- 1 file changed, 10 insertions(+), 7 deletions(-) diff --git a/folder_paths.py b/folder_paths.py index f1bf40f8c..a84524d9f 100644 --- a/folder_paths.py +++ b/folder_paths.py @@ -1,5 +1,6 @@ import os import time +import logging supported_pt_extensions = set(['.ckpt', '.pt', '.bin', '.pth', '.safetensors']) @@ -44,7 +45,7 @@ if not os.path.exists(input_directory): try: os.makedirs(input_directory) except: - print("Failed to create input directory") + logging.error("Failed to create input directory") def set_output_directory(output_dir): global output_directory @@ -146,21 +147,23 @@ def recursive_search(directory, excluded_dir_names=None): try: dirs[directory] = os.path.getmtime(directory) except FileNotFoundError: - print(f"Warning: Unable to access {directory}. Skipping this path.") - + logging.warning(f"Warning: Unable to access {directory}. Skipping this path.") + + logging.debug("recursive file list on directory {}".format(directory)) for dirpath, subdirs, filenames in os.walk(directory, followlinks=True, topdown=True): subdirs[:] = [d for d in subdirs if d not in excluded_dir_names] for file_name in filenames: relative_path = os.path.relpath(os.path.join(dirpath, file_name), directory) result.append(relative_path) - + for d in subdirs: path = os.path.join(dirpath, d) try: dirs[path] = os.path.getmtime(path) except FileNotFoundError: - print(f"Warning: Unable to access {path}. Skipping this path.") + logging.warning(f"Warning: Unable to access {path}. Skipping this path.") continue + logging.debug("found {} files".format(len(result))) return result, dirs def filter_files_extensions(files, extensions): @@ -248,8 +251,8 @@ def get_save_image_path(filename_prefix, output_dir, image_width=0, image_height err = "**** ERROR: Saving image outside the output folder is not allowed." + \ "\n full_output_folder: " + os.path.abspath(full_output_folder) + \ "\n output_dir: " + output_dir + \ - "\n commonpath: " + os.path.commonpath((output_dir, os.path.abspath(full_output_folder))) - print(err) + "\n commonpath: " + os.path.commonpath((output_dir, os.path.abspath(full_output_folder))) + logging.error(err) raise Exception(err) try: From 94a5a67c32ded77efbd1b47e8d0257ddab9ac3bc Mon Sep 17 00:00:00 2001 From: comfyanonymous Date: Fri, 29 Mar 2024 14:43:24 -0400 Subject: [PATCH 26/54] 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 27/54] 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 28/54] 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 29/54] --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 30/54] 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 31/54] 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 32/54] 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 33/54] 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 34/54] 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 35/54] 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 36/54] 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 37/54] 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 38/54] 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 39/54] 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 40/54] 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 41/54] 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 42/54] 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 43/54] 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 44/54] 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 45/54] 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 46/54] 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 47/54] 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 48/54] 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 49/54] 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 50/54] 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 51/54] 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 52/54] 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 53/54] 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 54/54] 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):