diff --git a/comfy/cldm/mmdit.py b/comfy/cldm/mmdit.py index 025c2fb5d..54a58ab83 100644 --- a/comfy/cldm/mmdit.py +++ b/comfy/cldm/mmdit.py @@ -6,6 +6,7 @@ class ControlNet(comfy.ldm.modules.diffusionmodules.mmdit.MMDiT): def __init__( self, num_blocks = None, + control_latent_channels = None, dtype = None, device = None, operations = None, @@ -17,10 +18,13 @@ class ControlNet(comfy.ldm.modules.diffusionmodules.mmdit.MMDiT): for _ in range(len(self.joint_blocks)): self.controlnet_blocks.append(operations.Linear(self.hidden_size, self.hidden_size, device=device, dtype=dtype)) + if control_latent_channels is None: + control_latent_channels = self.in_channels + self.pos_embed_input = comfy.ldm.modules.diffusionmodules.mmdit.PatchEmbed( None, self.patch_size, - self.in_channels, + control_latent_channels, self.hidden_size, bias=True, strict_img_size=False, diff --git a/comfy/controlnet.py b/comfy/controlnet.py index 860891965..1ea00eccf 100644 --- a/comfy/controlnet.py +++ b/comfy/controlnet.py @@ -79,13 +79,19 @@ class ControlBase: self.previous_controlnet = None self.extra_conds = [] self.strength_type = StrengthType.CONSTANT + self.concat_mask = False + self.extra_concat_orig = [] + self.extra_concat = None - def set_cond_hint(self, cond_hint, strength=1.0, timestep_percent_range=(0.0, 1.0), vae=None): + def set_cond_hint(self, cond_hint, strength=1.0, timestep_percent_range=(0.0, 1.0), vae=None, extra_concat=[]): self.cond_hint_original = cond_hint self.strength = strength self.timestep_percent_range = timestep_percent_range if self.latent_format is not None: self.vae = vae + self.extra_concat_orig = extra_concat.copy() + if self.concat_mask and len(self.extra_concat_orig) == 0: + self.extra_concat_orig.append(torch.tensor([[[[1.0]]]])) return self def pre_run(self, model, percent_to_timestep_function): @@ -100,9 +106,9 @@ class ControlBase: def cleanup(self): if self.previous_controlnet is not None: self.previous_controlnet.cleanup() - if self.cond_hint is not None: - del self.cond_hint - self.cond_hint = None + + self.cond_hint = None + self.extra_concat = None self.timestep_range = None def get_models(self): @@ -123,6 +129,8 @@ class ControlBase: c.vae = self.vae c.extra_conds = self.extra_conds.copy() c.strength_type = self.strength_type + c.concat_mask = self.concat_mask + c.extra_concat_orig = self.extra_concat_orig.copy() def inference_memory_requirements(self, dtype): if self.previous_controlnet is not None: @@ -175,7 +183,7 @@ class ControlBase: class ControlNet(ControlBase): - def __init__(self, control_model=None, global_average_pooling=False, compression_ratio=8, latent_format=None, device=None, load_device=None, manual_cast_dtype=None, extra_conds=["y"], strength_type=StrengthType.CONSTANT): + def __init__(self, control_model=None, global_average_pooling=False, compression_ratio=8, latent_format=None, device=None, load_device=None, manual_cast_dtype=None, extra_conds=["y"], strength_type=StrengthType.CONSTANT, concat_mask=False): super().__init__(device) self.control_model = control_model self.load_device = load_device @@ -189,6 +197,7 @@ class ControlNet(ControlBase): self.latent_format = latent_format self.extra_conds += extra_conds self.strength_type = strength_type + self.concat_mask = concat_mask def get_control(self, x_noisy, t, cond, batched_number): control_prev = None @@ -220,6 +229,13 @@ class ControlNet(ControlBase): comfy.model_management.load_models_gpu(loaded_models) if self.latent_format is not None: self.cond_hint = self.latent_format.process_in(self.cond_hint) + if len(self.extra_concat_orig) > 0: + to_concat = [] + for c in self.extra_concat_orig: + c = comfy.utils.common_upscale(c, self.cond_hint.shape[3], self.cond_hint.shape[2], self.upscale_algorithm, "center") + to_concat.append(comfy.utils.repeat_to_batch_size(c, self.cond_hint.shape[0])) + self.cond_hint = torch.cat([self.cond_hint] + to_concat, dim=1) + self.cond_hint = self.cond_hint.to(device=self.device, dtype=dtype) if x_noisy.shape[0] != self.cond_hint.shape[0]: self.cond_hint = broadcast_image_to(self.cond_hint, x_noisy.shape[0], batched_number) @@ -410,12 +426,17 @@ def load_controlnet_mmdit(sd): for k in sd: new_sd[k] = sd[k] - control_model = comfy.cldm.mmdit.ControlNet(num_blocks=num_blocks, operations=operations, device=offload_device, dtype=unet_dtype, **model_config.unet_config) + concat_mask = False + control_latent_channels = new_sd.get("pos_embed_input.proj.weight").shape[1] + if control_latent_channels == 17: #inpaint controlnet + concat_mask = True + + control_model = comfy.cldm.mmdit.ControlNet(num_blocks=num_blocks, control_latent_channels=control_latent_channels, operations=operations, device=offload_device, dtype=unet_dtype, **model_config.unet_config) control_model = controlnet_load_state_dict(control_model, new_sd) latent_format = comfy.latent_formats.SD3() latent_format.shift_factor = 0 #SD3 controlnet weirdness - control = ControlNet(control_model, compression_ratio=1, latent_format=latent_format, load_device=load_device, manual_cast_dtype=manual_cast_dtype) + control = ControlNet(control_model, compression_ratio=1, latent_format=latent_format, concat_mask=concat_mask, load_device=load_device, manual_cast_dtype=manual_cast_dtype) return control @@ -450,13 +471,16 @@ def load_controlnet_flux_instantx(sd): num_union_modes = new_sd[union_cnet].shape[0] control_latent_channels = new_sd.get("pos_embed_input.weight").shape[1] // 4 + concat_mask = False + if control_latent_channels == 17: + concat_mask = True control_model = comfy.ldm.flux.controlnet.ControlNetFlux(latent_input=True, num_union_modes=num_union_modes, control_latent_channels=control_latent_channels, operations=operations, device=offload_device, dtype=unet_dtype, **model_config.unet_config) control_model = controlnet_load_state_dict(control_model, new_sd) latent_format = comfy.latent_formats.Flux() extra_conds = ['y', 'guidance'] - control = ControlNet(control_model, compression_ratio=1, latent_format=latent_format, load_device=load_device, manual_cast_dtype=manual_cast_dtype, extra_conds=extra_conds) + control = ControlNet(control_model, compression_ratio=1, latent_format=latent_format, concat_mask=concat_mask, load_device=load_device, manual_cast_dtype=manual_cast_dtype, extra_conds=extra_conds) return control def convert_mistoline(sd): diff --git a/comfy/model_patcher.py b/comfy/model_patcher.py index 808601c72..0d4360977 100644 --- a/comfy/model_patcher.py +++ b/comfy/model_patcher.py @@ -325,17 +325,21 @@ class ModelPatcher: return list(p) def get_key_patches(self, filter_prefix=None): - comfy.model_management.unload_model_clones(self) model_sd = self.model_state_dict() p = {} for k in model_sd: if filter_prefix is not None: if not k.startswith(filter_prefix): continue - if k in self.patches: - p[k] = [model_sd[k]] + self.patches[k] + bk = self.backup.get(k, None) + if bk is not None: + weight = bk.weight else: - p[k] = (model_sd[k],) + weight = model_sd[k] + if k in self.patches: + p[k] = [weight] + self.patches[k] + else: + p[k] = (weight,) return p def model_state_dict(self, filter_prefix=None): diff --git a/comfy/sd.py b/comfy/sd.py index 1e1a6594b..07310b9d4 100644 --- a/comfy/sd.py +++ b/comfy/sd.py @@ -445,12 +445,8 @@ def load_text_encoder_state_dicts(state_dicts=[], embedding_directory=None, clip clip_target.tokenizer = comfy.text_encoders.sa_t5.SAT5Tokenizer else: w = clip_data[0].get("text_model.embeddings.position_embedding.weight", None) - if w is not None and w.shape[0] == 248: - clip_target.clip = comfy.text_encoders.long_clipl.LongClipModel - clip_target.tokenizer = comfy.text_encoders.long_clipl.LongClipTokenizer - else: - clip_target.clip = sd1_clip.SD1ClipModel - clip_target.tokenizer = sd1_clip.SD1Tokenizer + clip_target.clip = sd1_clip.SD1ClipModel + clip_target.tokenizer = sd1_clip.SD1Tokenizer elif len(clip_data) == 2: if clip_type == CLIPType.SD3: clip_target.clip = comfy.text_encoders.sd3_clip.sd3_clip(clip_l=True, clip_g=True, t5=False) @@ -475,10 +471,12 @@ def load_text_encoder_state_dicts(state_dicts=[], embedding_directory=None, clip clip_target.tokenizer = comfy.text_encoders.sd3_clip.SD3Tokenizer parameters = 0 + tokenizer_data = {} for c in clip_data: parameters += comfy.utils.calculate_parameters(c) + tokenizer_data, model_options = comfy.text_encoders.long_clipl.model_options_long_clip(c, tokenizer_data, model_options) - clip = CLIP(clip_target, embedding_directory=embedding_directory, parameters=parameters, model_options=model_options) + clip = CLIP(clip_target, embedding_directory=embedding_directory, parameters=parameters, tokenizer_data=tokenizer_data, model_options=model_options) for c in clip_data: m, u = clip.load_sd(c) if len(m) > 0: diff --git a/comfy/sd1_clip.py b/comfy/sd1_clip.py index 676653f77..9f0d95b0a 100644 --- a/comfy/sd1_clip.py +++ b/comfy/sd1_clip.py @@ -542,6 +542,7 @@ class SD1Tokenizer: def __init__(self, embedding_directory=None, tokenizer_data={}, clip_name="l", tokenizer=SDTokenizer): self.clip_name = clip_name self.clip = "clip_{}".format(self.clip_name) + tokenizer = tokenizer_data.get("{}_tokenizer_class".format(self.clip), tokenizer) setattr(self, self.clip, tokenizer(embedding_directory=embedding_directory, tokenizer_data=tokenizer_data)) def tokenize_with_weights(self, text:str, return_word_ids=False): @@ -570,6 +571,7 @@ class SD1ClipModel(torch.nn.Module): self.clip_name = clip_name self.clip = "clip_{}".format(self.clip_name) + clip_model = model_options.get("{}_class".format(self.clip), clip_model) setattr(self, self.clip, clip_model(device=device, dtype=dtype, model_options=model_options, **kwargs)) self.dtypes = set() diff --git a/comfy/sdxl_clip.py b/comfy/sdxl_clip.py index a0145caa4..4d0a4e8e7 100644 --- a/comfy/sdxl_clip.py +++ b/comfy/sdxl_clip.py @@ -22,7 +22,8 @@ class SDXLClipGTokenizer(sd1_clip.SDTokenizer): class SDXLTokenizer: def __init__(self, embedding_directory=None, tokenizer_data={}): - self.clip_l = sd1_clip.SDTokenizer(embedding_directory=embedding_directory) + clip_l_tokenizer_class = tokenizer_data.get("clip_l_tokenizer_class", sd1_clip.SDTokenizer) + self.clip_l = clip_l_tokenizer_class(embedding_directory=embedding_directory) self.clip_g = SDXLClipGTokenizer(embedding_directory=embedding_directory) def tokenize_with_weights(self, text:str, return_word_ids=False): @@ -40,7 +41,8 @@ class SDXLTokenizer: class SDXLClipModel(torch.nn.Module): def __init__(self, device="cpu", dtype=None, model_options={}): super().__init__() - self.clip_l = sd1_clip.SDClipModel(layer="hidden", layer_idx=-2, device=device, dtype=dtype, layer_norm_hidden_state=False, model_options=model_options) + clip_l_class = model_options.get("clip_l_class", sd1_clip.SDClipModel) + self.clip_l = clip_l_class(layer="hidden", layer_idx=-2, device=device, dtype=dtype, layer_norm_hidden_state=False, model_options=model_options) self.clip_g = SDXLClipG(device=device, dtype=dtype, model_options=model_options) self.dtypes = set([dtype]) @@ -57,7 +59,8 @@ class SDXLClipModel(torch.nn.Module): token_weight_pairs_l = token_weight_pairs["l"] g_out, g_pooled = self.clip_g.encode_token_weights(token_weight_pairs_g) l_out, l_pooled = self.clip_l.encode_token_weights(token_weight_pairs_l) - return torch.cat([l_out, g_out], dim=-1), g_pooled + cut_to = min(l_out.shape[1], g_out.shape[1]) + return torch.cat([l_out[:,:cut_to], g_out[:,:cut_to]], dim=-1), g_pooled def load_sd(self, sd): if "text_model.encoder.layers.30.mlp.fc1.weight" in sd: diff --git a/comfy/text_encoders/flux.py b/comfy/text_encoders/flux.py index 91d1f249d..7c75fe64b 100644 --- a/comfy/text_encoders/flux.py +++ b/comfy/text_encoders/flux.py @@ -18,7 +18,8 @@ class T5XXLTokenizer(sd1_clip.SDTokenizer): class FluxTokenizer: def __init__(self, embedding_directory=None, tokenizer_data={}): - self.clip_l = sd1_clip.SDTokenizer(embedding_directory=embedding_directory) + clip_l_tokenizer_class = tokenizer_data.get("clip_l_tokenizer_class", sd1_clip.SDTokenizer) + self.clip_l = clip_l_tokenizer_class(embedding_directory=embedding_directory) self.t5xxl = T5XXLTokenizer(embedding_directory=embedding_directory) def tokenize_with_weights(self, text:str, return_word_ids=False): @@ -38,7 +39,8 @@ class FluxClipModel(torch.nn.Module): def __init__(self, dtype_t5=None, device="cpu", dtype=None, model_options={}): super().__init__() dtype_t5 = comfy.model_management.pick_weight_dtype(dtype_t5, dtype, device) - self.clip_l = sd1_clip.SDClipModel(device=device, dtype=dtype, return_projected_pooled=False, model_options=model_options) + clip_l_class = model_options.get("clip_l_class", sd1_clip.SDClipModel) + self.clip_l = clip_l_class(device=device, dtype=dtype, return_projected_pooled=False, model_options=model_options) self.t5xxl = T5XXLModel(device=device, dtype=dtype_t5, model_options=model_options) self.dtypes = set([dtype, dtype_t5]) diff --git a/comfy/text_encoders/long_clipl.py b/comfy/text_encoders/long_clipl.py index 4677fb3b0..b81912cb3 100644 --- a/comfy/text_encoders/long_clipl.py +++ b/comfy/text_encoders/long_clipl.py @@ -6,9 +6,9 @@ class LongClipTokenizer_(sd1_clip.SDTokenizer): super().__init__(max_length=248, embedding_directory=embedding_directory, tokenizer_data=tokenizer_data) class LongClipModel_(sd1_clip.SDClipModel): - def __init__(self, device="cpu", dtype=None, model_options={}): + def __init__(self, *args, **kwargs): textmodel_json_config = os.path.join(os.path.dirname(os.path.realpath(__file__)), "long_clipl.json") - super().__init__(device=device, textmodel_json_config=textmodel_json_config, return_projected_pooled=False, dtype=dtype, model_options=model_options) + super().__init__(*args, textmodel_json_config=textmodel_json_config, **kwargs) class LongClipTokenizer(sd1_clip.SD1Tokenizer): def __init__(self, embedding_directory=None, tokenizer_data={}): @@ -17,3 +17,14 @@ class LongClipTokenizer(sd1_clip.SD1Tokenizer): class LongClipModel(sd1_clip.SD1ClipModel): def __init__(self, device="cpu", dtype=None, model_options={}, **kwargs): super().__init__(device=device, dtype=dtype, model_options=model_options, clip_model=LongClipModel_, **kwargs) + +def model_options_long_clip(sd, tokenizer_data, model_options): + w = sd.get("clip_l.text_model.embeddings.position_embedding.weight", None) + if w is None: + w = sd.get("text_model.embeddings.position_embedding.weight", None) + if w is not None and w.shape[0] == 248: + tokenizer_data = tokenizer_data.copy() + model_options = model_options.copy() + tokenizer_data["clip_l_tokenizer_class"] = LongClipTokenizer_ + model_options["clip_l_class"] = LongClipModel_ + return tokenizer_data, model_options diff --git a/comfy/text_encoders/sd3_clip.py b/comfy/text_encoders/sd3_clip.py index e3832ac24..c54f28856 100644 --- a/comfy/text_encoders/sd3_clip.py +++ b/comfy/text_encoders/sd3_clip.py @@ -20,7 +20,8 @@ class T5XXLTokenizer(sd1_clip.SDTokenizer): class SD3Tokenizer: def __init__(self, embedding_directory=None, tokenizer_data={}): - self.clip_l = sd1_clip.SDTokenizer(embedding_directory=embedding_directory) + clip_l_tokenizer_class = tokenizer_data.get("clip_l_tokenizer_class", sd1_clip.SDTokenizer) + self.clip_l = clip_l_tokenizer_class(embedding_directory=embedding_directory) self.clip_g = sdxl_clip.SDXLClipGTokenizer(embedding_directory=embedding_directory) self.t5xxl = T5XXLTokenizer(embedding_directory=embedding_directory) @@ -42,7 +43,8 @@ class SD3ClipModel(torch.nn.Module): super().__init__() self.dtypes = set() if clip_l: - self.clip_l = sd1_clip.SDClipModel(layer="hidden", layer_idx=-2, device=device, dtype=dtype, layer_norm_hidden_state=False, return_projected_pooled=False, model_options=model_options) + clip_l_class = model_options.get("clip_l_class", sd1_clip.SDClipModel) + self.clip_l = clip_l_class(layer="hidden", layer_idx=-2, device=device, dtype=dtype, layer_norm_hidden_state=False, return_projected_pooled=False, model_options=model_options) self.dtypes.add(dtype) else: self.clip_l = None @@ -95,7 +97,8 @@ class SD3ClipModel(torch.nn.Module): if self.clip_g is not None: g_out, g_pooled = self.clip_g.encode_token_weights(token_weight_pairs_g) if lg_out is not None: - lg_out = torch.cat([lg_out, g_out], dim=-1) + cut_to = min(lg_out.shape[1], g_out.shape[1]) + lg_out = torch.cat([lg_out[:,:cut_to], g_out[:,:cut_to]], dim=-1) else: lg_out = torch.nn.functional.pad(g_out, (768, 0)) else: diff --git a/comfy_execution/caching.py b/comfy_execution/caching.py index 62311ed7f..630f280fc 100644 --- a/comfy_execution/caching.py +++ b/comfy_execution/caching.py @@ -1,11 +1,21 @@ import itertools -from typing import Sequence, Mapping +from typing import Sequence, Mapping, Dict from comfy_execution.graph import DynamicPrompt import nodes from comfy_execution.graph_utils import is_link +NODE_CLASS_CONTAINS_UNIQUE_ID: Dict[str, bool] = {} + + +def include_unique_id_in_input(class_type: str) -> bool: + if class_type in NODE_CLASS_CONTAINS_UNIQUE_ID: + return NODE_CLASS_CONTAINS_UNIQUE_ID[class_type] + class_def = nodes.NODE_CLASS_MAPPINGS[class_type] + NODE_CLASS_CONTAINS_UNIQUE_ID[class_type] = "UNIQUE_ID" in class_def.INPUT_TYPES().get("hidden", {}).values() + return NODE_CLASS_CONTAINS_UNIQUE_ID[class_type] + class CacheKeySet: def __init__(self, dynprompt, node_ids, is_changed_cache): self.keys = {} @@ -98,7 +108,7 @@ class CacheKeySetInputSignature(CacheKeySet): class_type = node["class_type"] class_def = nodes.NODE_CLASS_MAPPINGS[class_type] signature = [class_type, self.is_changed_cache.get(node_id)] - if self.include_node_id_in_input() or (hasattr(class_def, "NOT_IDEMPOTENT") and class_def.NOT_IDEMPOTENT) or "UNIQUE_ID" in class_def.INPUT_TYPES().get("hidden", {}).values(): + if self.include_node_id_in_input() or (hasattr(class_def, "NOT_IDEMPOTENT") and class_def.NOT_IDEMPOTENT) or include_unique_id_in_input(class_type): signature.append(node_id) inputs = node["inputs"] for key in sorted(inputs.keys()): diff --git a/comfy_extras/nodes_audio.py b/comfy_extras/nodes_audio.py index 6990b3f9a..1c763f259 100644 --- a/comfy_extras/nodes_audio.py +++ b/comfy_extras/nodes_audio.py @@ -58,6 +58,9 @@ class VAEDecodeAudio: def decode(self, vae, samples): audio = vae.decode(samples["samples"]).movedim(-1, 1) + std = torch.std(audio, dim=[1,2], keepdim=True) * 5.0 + std[std < 1.0] = 1.0 + audio /= std return ({"waveform": audio, "sample_rate": 44100}, ) diff --git a/comfy_extras/nodes_controlnet.py b/comfy_extras/nodes_controlnet.py index 0773c8a5b..2d20e1fed 100644 --- a/comfy_extras/nodes_controlnet.py +++ b/comfy_extras/nodes_controlnet.py @@ -1,4 +1,6 @@ from comfy.cldm.control_types import UNION_CONTROLNET_TYPES +import nodes +import comfy.utils class SetUnionControlNetType: @classmethod @@ -22,6 +24,37 @@ class SetUnionControlNetType: return (control_net,) +class ControlNetInpaintingAliMamaApply(nodes.ControlNetApplyAdvanced): + @classmethod + def INPUT_TYPES(s): + return {"required": {"positive": ("CONDITIONING", ), + "negative": ("CONDITIONING", ), + "control_net": ("CONTROL_NET", ), + "vae": ("VAE", ), + "image": ("IMAGE", ), + "mask": ("MASK", ), + "strength": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 10.0, "step": 0.01}), + "start_percent": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 1.0, "step": 0.001}), + "end_percent": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.001}) + }} + + FUNCTION = "apply_inpaint_controlnet" + + CATEGORY = "conditioning/controlnet" + + def apply_inpaint_controlnet(self, positive, negative, control_net, vae, image, mask, strength, start_percent, end_percent): + extra_concat = [] + if control_net.concat_mask: + mask = 1.0 - mask.reshape((-1, 1, mask.shape[-2], mask.shape[-1])) + mask_apply = comfy.utils.common_upscale(mask, image.shape[2], image.shape[1], "bilinear", "center").round() + image = image * mask_apply.movedim(1, -1).repeat(1, 1, 1, image.shape[3]) + extra_concat = [mask] + + return self.apply_controlnet(positive, negative, control_net, image, strength, start_percent, end_percent, vae=vae, extra_concat=extra_concat) + + + NODE_CLASS_MAPPINGS = { "SetUnionControlNetType": SetUnionControlNetType, + "ControlNetInpaintingAliMamaApply": ControlNetInpaintingAliMamaApply, } diff --git a/nodes.py b/nodes.py index 36022f731..70f52741e 100644 --- a/nodes.py +++ b/nodes.py @@ -824,7 +824,7 @@ class ControlNetApplyAdvanced: CATEGORY = "conditioning/controlnet" - def apply_controlnet(self, positive, negative, control_net, image, strength, start_percent, end_percent, vae=None): + def apply_controlnet(self, positive, negative, control_net, image, strength, start_percent, end_percent, vae=None, extra_concat=[]): if strength == 0: return (positive, negative) @@ -841,7 +841,7 @@ class ControlNetApplyAdvanced: if prev_cnet in cnets: c_net = cnets[prev_cnet] else: - c_net = control_net.copy().set_cond_hint(control_hint, strength, (start_percent, end_percent), vae) + c_net = control_net.copy().set_cond_hint(control_hint, strength, (start_percent, end_percent), vae=vae, extra_concat=extra_concat) c_net.set_previous_controlnet(prev_cnet) cnets[prev_cnet] = c_net