From c808e875c4d293700128981e80a76d5f0bf30025 Mon Sep 17 00:00:00 2001 From: Jedrzej Kosinski Date: Thu, 23 May 2024 15:59:50 -0500 Subject: [PATCH 1/9] Implement fallback_attention_mm for SparseCtrl to mirror AnimateDiff-Evolved --- adv_control/control_sparsectrl.py | 33 ++++++++++++++++++++++++++++--- 1 file changed, 30 insertions(+), 3 deletions(-) diff --git a/adv_control/control_sparsectrl.py b/adv_control/control_sparsectrl.py index 60ffbcc..23a49c8 100644 --- a/adv_control/control_sparsectrl.py +++ b/adv_control/control_sparsectrl.py @@ -3,6 +3,7 @@ #and then taken from comfy/cldm/cldm.py and modified again from abc import ABC, abstractmethod +import copy import math import numpy as np from typing import Iterable, Union @@ -34,12 +35,18 @@ from .utils import TimestepKeyframeGroup, disable_weight_init_clean_groupnorm, p # until xformers bug is fixed, do not use xformers for VersatileAttention! TODO: change this when fix is out # logic for choosing optimized_attention method taken from comfy/ldm/modules/attention.py +# a fallback_attention_mm is selected to avoid CUDA configuration limitation with pytorch's scaled_dot_product optimized_attention_mm = attention_basic +fallback_attention_mm = attention_basic if comfy.model_management.xformers_enabled(): pass #optimized_attention_mm = attention_xformers if comfy.model_management.pytorch_attention_enabled(): optimized_attention_mm = attention_pytorch + if args.use_split_cross_attention: + fallback_attention_mm = attention_split + else: + fallback_attention_mm = attention_sub_quad else: if args.use_split_cross_attention: optimized_attention_mm = attention_split @@ -641,6 +648,8 @@ class TemporalTransformer3DModel(nn.Module): del self.temp_scale_mask self.temp_scale_mask = None self.prev_hidden_states_batch = 0 + for block in self.transformer_blocks: + block.reset_temp_vars() def get_scale_mask(self, hidden_states: Tensor) -> Union[Tensor, None]: # if no raw mask, return None @@ -779,6 +788,10 @@ class TemporalTransformerBlock(nn.Module): for block in self.attention_blocks: block.set_sub_idxs(sub_idxs) + def reset_temp_vars(self): + for block in self.attention_blocks: + block.reset_temp_vars() + def forward( self, hidden_states, @@ -833,13 +846,14 @@ class PositionalEncoding(nn.Module): return self.dropout(x) -class CrossAttentionMM(nn.Module): +class CrossAttentionMMSparse(nn.Module): def __init__(self, query_dim, context_dim=None, heads=8, dim_head=64, dropout=0., dtype=None, device=None, operations=comfy.ops.disable_weight_init): super().__init__() inner_dim = dim_head * heads context_dim = default(context_dim, query_dim) + self.actual_attention = optimized_attention_mm self.heads = heads self.dim_head = dim_head self.scale = None @@ -850,6 +864,9 @@ class CrossAttentionMM(nn.Module): self.to_out = nn.Sequential(operations.Linear(inner_dim, query_dim, dtype=dtype, device=device), nn.Dropout(dropout)) + def reset_attention_type(self): + self.actual_attention = optimized_attention_mm + def forward(self, x, context=None, value=None, mask=None, scale_mask=None): q = self.to_q(x) context = default(context, x) @@ -868,11 +885,18 @@ class CrossAttentionMM(nn.Module): if scale_mask is not None: k *= scale_mask - out = optimized_attention_mm(q, k, v, self.heads, mask) + try: + out = self.actual_attention(q, k, v, self.heads, mask) + except RuntimeError as e: + if str(e).startswith("CUDA error: invalid configuration argument"): + self.actual_attention = fallback_attention_mm + out = self.actual_attention(q, k, v, self.heads, mask) + else: + raise return self.to_out(out) -class VersatileAttention(CrossAttentionMM): +class VersatileAttention(CrossAttentionMMSparse): def __init__( self, attention_mode=None, @@ -911,6 +935,9 @@ class VersatileAttention(CrossAttentionMM): if self.pos_encoder != None: self.pos_encoder.set_sub_idxs(sub_idxs) + def reset_temp_vars(self): + self.reset_attention_type() + def forward( self, hidden_states: Tensor, From 5660a110df7cff8740ac5ad57d10c192d318c434 Mon Sep 17 00:00:00 2001 From: Jedrzej Kosinski Date: Thu, 23 May 2024 18:12:28 -0500 Subject: [PATCH 2/9] Made SparseCtrl work with fp8 and cleaner lowvram support --- adv_control/control.py | 13 +++-- adv_control/control_sparsectrl.py | 86 +++++++++++++++++-------------- 2 files changed, 54 insertions(+), 45 deletions(-) diff --git a/adv_control/control.py b/adv_control/control.py index bea5aa7..b9253de 100644 --- a/adv_control/control.py +++ b/adv_control/control.py @@ -3,6 +3,7 @@ from torch import Tensor import torch import os +import comfy.ops import comfy.utils import comfy.model_management import comfy.model_detection @@ -551,13 +552,9 @@ def load_sparsectrl(ckpt_path: str, controlnet_data: dict[str, Tensor]=None, tim motion_data[key] = controlnet_data.pop(key) if len(motion_data) == 0: raise ValueError(f"No motion-related keys in '{ckpt_path}'; not a valid SparseCtrl model!") - motion_wrapper: SparseCtrlMotionWrapper = SparseCtrlMotionWrapper(motion_data).to(comfy.model_management.unet_dtype()) - missing, unexpected = motion_wrapper.load_state_dict(motion_data) - if len(missing) > 0 or len(unexpected) > 0: - logger.info(f"SparseCtrlMotionWrapper: {missing}, {unexpected}") # now, load as if it was a normal controlnet - mostly copied from comfy load_controlnet function - controlnet_config = None + controlnet_config: dict[str] = None is_diffusers = False use_simplified_conditioning_embedding = False if "controlnet_cond_embedding.conv_in.weight" in controlnet_data: @@ -681,6 +678,12 @@ def load_sparsectrl(ckpt_path: str, controlnet_data: dict[str, Tensor]=None, tim if filename.endswith("_shuffle") or filename.endswith("_shuffle_fp16"): #TODO: smarter way of enabling global_average_pooling global_average_pooling = True + # actually load motion portion of model now + motion_wrapper: SparseCtrlMotionWrapper = SparseCtrlMotionWrapper(motion_data, ops=controlnet_config.get("operations", None)).to(comfy.model_management.unet_dtype()) + missing, unexpected = motion_wrapper.load_state_dict(motion_data) + if len(missing) > 0 or len(unexpected) > 0: + logger.info(f"SparseCtrlMotionWrapper: {missing}, {unexpected}") + # both motion portion and controlnet portions are loaded; bring them together if using motion model if sparse_settings.use_motion: motion_wrapper.inject(control_model) diff --git a/adv_control/control_sparsectrl.py b/adv_control/control_sparsectrl.py index 23a49c8..58b4d85 100644 --- a/adv_control/control_sparsectrl.py +++ b/adv_control/control_sparsectrl.py @@ -107,28 +107,27 @@ class SparseModelPatcher(ModelPatcher): self.model: SparseControlNet super().__init__(*args, **kwargs) - def patch_model(self, device_to=None, patch_weights=True): - if patch_weights: - patched_model = super().patch_model(device_to) - else: - patched_model = super().patch_model(device_to, patch_weights) - try: - if self.model.motion_wrapper is not None: - self.model.motion_wrapper.to(device=device_to) - except Exception: - pass - return patched_model + def patch_model_lowvram(self, device_to=None, *args, **kwargs): + patched_model = super().patch_model_lowvram(device_to, *args, **kwargs) - def unpatch_model(self, device_to=None, unpatch_weights=True): - try: - if self.model.motion_wrapper is not None: - self.model.motion_wrapper.to(device=device_to) - except Exception: - pass - if unpatch_weights: - return super().unpatch_model(device_to) - else: - return super().unpatch_model(device_to, unpatch_weights) + if self.model.motion_wrapper is not None: + # figure out the tensors (likely pe's) that should be cast to device besides just the named_modules + remaining_tensors = list(self.model.motion_wrapper.state_dict().keys()) + named_modules = [] + for n, _ in self.model.motion_wrapper.named_modules(): + named_modules.append(n) + named_modules.append(f"{n}.weight") + named_modules.append(f"{n}.bias") + for name in named_modules: + if name in remaining_tensors: + remaining_tensors.remove(name) + + for key in remaining_tensors: + self.patch_weight_to_device(key, device_to) + if device_to is not None: + comfy.utils.set_attr(self.model.motion_wrapper, key, comfy.utils.get_attr(self.model.motion_wrapper, key).to(device_to)) + + return patched_model def clone(self): # normal ModelPatcher clone actions @@ -325,7 +324,7 @@ def get_position_encoding_max_len(mm_state_dict: dict[str, Tensor], mm_name: str class SparseCtrlMotionWrapper(nn.Module): - def __init__(self, mm_state_dict: dict[str, Tensor]): + def __init__(self, mm_state_dict: dict[str, Tensor], ops=disable_weight_init_clean_groupnorm): super().__init__() self.down_blocks: Iterable[MotionModule] = None self.up_blocks: Iterable[MotionModule] = None @@ -335,13 +334,13 @@ class SparseCtrlMotionWrapper(nn.Module): if get_down_block_max(mm_state_dict) > -1: self.down_blocks = nn.ModuleList([]) for c in layer_channels: - self.down_blocks.append(MotionModule(c, temporal_position_encoding_max_len=self.encoding_max_len, block_type=BlockType.DOWN)) + self.down_blocks.append(MotionModule(c, temporal_position_encoding_max_len=self.encoding_max_len, block_type=BlockType.DOWN, ops=ops)) if get_up_block_max(mm_state_dict) > -1: self.up_blocks = nn.ModuleList([]) for c in reversed(layer_channels): - self.up_blocks.append(MotionModule(c, temporal_position_encoding_max_len=self.encoding_max_len, block_type=BlockType.UP)) + self.up_blocks.append(MotionModule(c, temporal_position_encoding_max_len=self.encoding_max_len, block_type=BlockType.UP, ops=ops)) if has_mid_block(mm_state_dict): - self.mid_block = MotionModule(1280, temporal_position_encoding_max_len=self.encoding_max_len, block_type=BlockType.MID) + self.mid_block = MotionModule(1280, temporal_position_encoding_max_len=self.encoding_max_len, block_type=BlockType.MID, ops=ops) def inject(self, unet: SparseControlNet): # inject input (down) blocks @@ -455,22 +454,22 @@ class SparseCtrlMotionWrapper(nn.Module): class MotionModule(nn.Module): - def __init__(self, in_channels, temporal_position_encoding_max_len=24, block_type: str=BlockType.DOWN): + def __init__(self, in_channels, temporal_position_encoding_max_len=24, block_type: str=BlockType.DOWN, ops=disable_weight_init_clean_groupnorm): super().__init__() if block_type == BlockType.MID: # mid blocks contain only a single VanillaTemporalModule - self.motion_modules: Iterable[VanillaTemporalModule] = nn.ModuleList([get_motion_module(in_channels, temporal_position_encoding_max_len)]) + self.motion_modules: Iterable[VanillaTemporalModule] = nn.ModuleList([get_motion_module(in_channels, temporal_position_encoding_max_len, ops=ops)]) else: # down blocks contain two VanillaTemporalModules self.motion_modules: Iterable[VanillaTemporalModule] = nn.ModuleList( [ - get_motion_module(in_channels, temporal_position_encoding_max_len), - get_motion_module(in_channels, temporal_position_encoding_max_len) + get_motion_module(in_channels, temporal_position_encoding_max_len, ops=ops), + get_motion_module(in_channels, temporal_position_encoding_max_len, ops=ops) ] ) # up blocks contain one additional VanillaTemporalModule if block_type == BlockType.UP: - self.motion_modules.append(get_motion_module(in_channels, temporal_position_encoding_max_len)) + self.motion_modules.append(get_motion_module(in_channels, temporal_position_encoding_max_len, ops=ops)) def set_video_length(self, video_length: int, full_length: int): for motion_module in self.motion_modules: @@ -497,9 +496,9 @@ class MotionModule(nn.Module): motion_module.reset_temp_vars() -def get_motion_module(in_channels, temporal_position_encoding_max_len): +def get_motion_module(in_channels, temporal_position_encoding_max_len, ops=disable_weight_init_clean_groupnorm): # unlike normal AD, there is only one attention block expected in SparseCtrl models - return VanillaTemporalModule(in_channels=in_channels, attention_block_types=("Temporal_Self",), temporal_position_encoding_max_len=temporal_position_encoding_max_len) + return VanillaTemporalModule(in_channels=in_channels, attention_block_types=("Temporal_Self",), temporal_position_encoding_max_len=temporal_position_encoding_max_len, ops=ops) class VanillaTemporalModule(nn.Module): @@ -514,6 +513,7 @@ class VanillaTemporalModule(nn.Module): temporal_position_encoding_max_len=24, temporal_attention_dim_div=1, zero_initialize=True, + ops=disable_weight_init_clean_groupnorm, ): super().__init__() self.strength = 1.0 @@ -528,6 +528,7 @@ class VanillaTemporalModule(nn.Module): cross_frame_attention_mode=cross_frame_attention_mode, temporal_position_encoding=temporal_position_encoding, temporal_position_encoding_max_len=temporal_position_encoding_max_len, + ops=ops, ) if zero_initialize: @@ -585,6 +586,7 @@ class TemporalTransformer3DModel(nn.Module): cross_frame_attention_mode=None, temporal_position_encoding=False, temporal_position_encoding_max_len=24, + ops=disable_weight_init_clean_groupnorm, ): super().__init__() self.video_length = 16 @@ -599,10 +601,10 @@ class TemporalTransformer3DModel(nn.Module): inner_dim = num_attention_heads * attention_head_dim - self.norm = disable_weight_init_clean_groupnorm.GroupNorm( + self.norm = ops.GroupNorm( num_groups=norm_num_groups, num_channels=in_channels, eps=1e-6, affine=True ) - self.proj_in = nn.Linear(in_channels, inner_dim) + self.proj_in = ops.Linear(in_channels, inner_dim) self.transformer_blocks: Iterable[TemporalTransformerBlock] = nn.ModuleList( [ @@ -620,11 +622,12 @@ class TemporalTransformer3DModel(nn.Module): cross_frame_attention_mode=cross_frame_attention_mode, temporal_position_encoding=temporal_position_encoding, temporal_position_encoding_max_len=temporal_position_encoding_max_len, + ops=ops, ) for d in range(num_layers) ] ) - self.proj_out = nn.Linear(inner_dim, in_channels) + self.proj_out = ops.Linear(inner_dim, in_channels) def set_video_length(self, video_length: int, full_length: int): self.video_length = video_length @@ -748,6 +751,7 @@ class TemporalTransformerBlock(nn.Module): cross_frame_attention_mode=None, temporal_position_encoding=False, temporal_position_encoding_max_len=24, + ops=disable_weight_init_clean_groupnorm, ): super().__init__() @@ -770,15 +774,16 @@ class TemporalTransformerBlock(nn.Module): cross_frame_attention_mode=cross_frame_attention_mode, temporal_position_encoding=temporal_position_encoding, temporal_position_encoding_max_len=temporal_position_encoding_max_len, + ops=ops, ) ) - norms.append(nn.LayerNorm(dim)) + norms.append(ops.LayerNorm(dim)) self.attention_blocks: Iterable[VersatileAttention] = nn.ModuleList(attention_blocks) self.norms = nn.ModuleList(norms) - self.ff = FeedForward(dim, dropout=dropout, glu=(activation_fn == "geglu")) - self.ff_norm = nn.LayerNorm(dim) + self.ff = FeedForward(dim, dropout=dropout, glu=(activation_fn == "geglu"), operations=ops) + self.ff_norm = ops.LayerNorm(dim) def set_scale_multiplier(self, multiplier: Union[float, None]): for block in self.attention_blocks: @@ -848,7 +853,7 @@ class PositionalEncoding(nn.Module): class CrossAttentionMMSparse(nn.Module): def __init__(self, query_dim, context_dim=None, heads=8, dim_head=64, dropout=0., dtype=None, device=None, - operations=comfy.ops.disable_weight_init): + operations=disable_weight_init_clean_groupnorm): super().__init__() inner_dim = dim_head * heads context_dim = default(context_dim, query_dim) @@ -903,10 +908,11 @@ class VersatileAttention(CrossAttentionMMSparse): cross_frame_attention_mode=None, temporal_position_encoding=False, temporal_position_encoding_max_len=24, + ops=disable_weight_init_clean_groupnorm, *args, **kwargs, ): - super().__init__(*args, **kwargs) + super().__init__(operations=ops, *args, **kwargs) assert attention_mode == "Temporal" self.attention_mode = attention_mode From 1f95fac9648d0d40bf6a26053cf3a94cd8570aa2 Mon Sep 17 00:00:00 2001 From: Jedrzej Kosinski Date: Fri, 24 May 2024 00:11:26 -0500 Subject: [PATCH 3/9] Fixed odd behavior when context options are active and there are less condhints than latents, condhints/masks are extended (last image/mask) instead of repeated, starting to work on uncond_mask to extend uncond_multiplier feature, modified motion_strength > 1.0 behavior --- adv_control/control.py | 41 ++++++++++++++++++++----------- adv_control/control_reference.py | 4 +-- adv_control/control_sparsectrl.py | 13 +++++----- adv_control/utils.py | 23 ++++++++++++----- 4 files changed, 52 insertions(+), 29 deletions(-) diff --git a/adv_control/control.py b/adv_control/control.py index b9253de..235bb02 100644 --- a/adv_control/control.py +++ b/adv_control/control.py @@ -8,14 +8,15 @@ import comfy.utils import comfy.model_management import comfy.model_detection import comfy.controlnet as comfy_cn -from comfy.controlnet import ControlBase, ControlNet, ControlLora, T2IAdapter, broadcast_image_to +from comfy.controlnet import ControlBase, ControlNet, ControlLora, T2IAdapter from comfy.model_patcher import ModelPatcher from .control_sparsectrl import SparseModelPatcher, SparseControlNet, SparseCtrlMotionWrapper, SparseMethod, SparseSettings, SparseSpreadMethod, PreprocSparseRGBWrapper from .control_lllite import LLLiteModule, LLLitePatch from .control_svd import svd_unet_config_from_diffusers_unet, SVDControlNet, svd_unet_to_diffusers from .utils import (AdvancedControlBase, TimestepKeyframeGroup, LatentKeyframeGroup, ControlWeightType, ControlWeights, WeightTypeException, - manual_cast_clean_groupnorm, disable_weight_init_clean_groupnorm, prepare_mask_batch, get_properly_arranged_t2i_weights, load_torch_file_with_dict_factory) + manual_cast_clean_groupnorm, disable_weight_init_clean_groupnorm, prepare_mask_batch, get_properly_arranged_t2i_weights, load_torch_file_with_dict_factory, + broadcast_image_to_extend, extend_to_batch_size) from .logger import logger @@ -56,12 +57,15 @@ class ControlNetAdvanced(ControlNet, AdvancedControlBase): del self.cond_hint self.cond_hint = None # if self.cond_hint_original length greater or equal to real latent count, subdivide it before scaling - if self.sub_idxs is not None and self.cond_hint_original.size(0) >= self.full_latent_length: - self.cond_hint = comfy.utils.common_upscale(self.cond_hint_original[self.sub_idxs], x_noisy.shape[3] * 8, x_noisy.shape[2] * 8, 'nearest-exact', "center").to(dtype).to(self.device) + if self.sub_idxs is not None: + actual_cond_hint_orig = self.cond_hint_original + if self.cond_hint_original.size(0) < self.full_latent_length: + actual_cond_hint_orig = extend_to_batch_size(tensor=actual_cond_hint_orig, batch_size=self.full_latent_length) + self.cond_hint = comfy.utils.common_upscale(actual_cond_hint_orig[self.sub_idxs], x_noisy.shape[3] * 8, x_noisy.shape[2] * 8, 'nearest-exact', "center").to(dtype).to(self.device) else: self.cond_hint = comfy.utils.common_upscale(self.cond_hint_original, x_noisy.shape[3] * 8, x_noisy.shape[2] * 8, 'nearest-exact', "center").to(dtype).to(self.device) 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) + self.cond_hint = broadcast_image_to_extend(self.cond_hint, x_noisy.shape[0], batched_number) # prepare mask_cond_hint self.prepare_mask_cond_hint(x_noisy=x_noisy, t=t, cond=cond, batched_number=batched_number, dtype=dtype) @@ -98,7 +102,7 @@ class T2IAdapterAdvanced(T2IAdapter, AdvancedControlBase): def control_merge_inject(self, control_input, control_output, control_prev, output_dtype): # if has uncond multiplier, need to make sure control shapes are the same batch size as expected - if self.weights.has_uncond_multiplier: + if self.weights.has_uncond_multiplier or self.weights.has_uncond_mask: if control_input is not None: for i in range(len(control_input)): x = control_input[i] @@ -132,9 +136,12 @@ class T2IAdapterAdvanced(T2IAdapter, AdvancedControlBase): if self.sub_idxs is not None: # cond hints full_cond_hint_original = self.cond_hint_original + actual_cond_hint_orig = full_cond_hint_original del self.cond_hint self.cond_hint = None - self.cond_hint_original = full_cond_hint_original[self.sub_idxs] + if full_cond_hint_original.size(0) < self.full_latent_length: + actual_cond_hint_orig = extend_to_batch_size(tensor=full_cond_hint_original, batch_size=full_cond_hint_original.size(0)) + self.cond_hint_original = actual_cond_hint_orig[self.sub_idxs] # mask hints self.prepare_mask_cond_hint(x_noisy=x_noisy, t=t, cond=cond, batched_number=batched_number) return super().get_control(x_noisy, t, cond, batched_number) @@ -222,12 +229,15 @@ class SVDControlNetAdvanced(ControlNetAdvanced): del self.cond_hint self.cond_hint = None # if self.cond_hint_original length greater or equal to real latent count, subdivide it before scaling - if self.sub_idxs is not None and self.cond_hint_original.size(0) >= self.full_latent_length: - self.cond_hint = comfy.utils.common_upscale(self.cond_hint_original[self.sub_idxs], x_noisy.shape[3] * 8, x_noisy.shape[2] * 8, 'nearest-exact', "center").to(dtype).to(self.device) + if self.sub_idxs is not None: + actual_cond_hint_orig = self.cond_hint_original + if self.cond_hint_original.size(0) < self.full_latent_length: + actual_cond_hint_orig = extend_to_batch_size(tensor=actual_cond_hint_orig, batch_size=self.full_latent_length) + self.cond_hint = comfy.utils.common_upscale(actual_cond_hint_orig[self.sub_idxs], x_noisy.shape[3] * 8, x_noisy.shape[2] * 8, 'nearest-exact', "center").to(dtype).to(self.device) else: self.cond_hint = comfy.utils.common_upscale(self.cond_hint_original, x_noisy.shape[3] * 8, x_noisy.shape[2] * 8, 'nearest-exact', "center").to(dtype).to(self.device) 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) + self.cond_hint = broadcast_image_to_extend(self.cond_hint, x_noisy.shape[0], batched_number) # prepare mask_cond_hint self.prepare_mask_cond_hint(x_noisy=x_noisy, t=t, cond=cond, batched_number=batched_number, dtype=dtype) @@ -328,7 +338,7 @@ class SparseCtrlAdvanced(ControlNetAdvanced): del cond_mask # make cond_hint match x_noisy batch 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) + self.cond_hint = broadcast_image_to_extend(self.cond_hint, x_noisy.shape[0], batched_number) # prepare mask_cond_hint self.prepare_mask_cond_hint(x_noisy=x_noisy, t=t, cond=cond, batched_number=batched_number, dtype=dtype) @@ -412,12 +422,15 @@ class ControlLLLiteAdvanced(ControlBase, AdvancedControlBase): del self.cond_hint self.cond_hint = None # if self.cond_hint_original length greater or equal to real latent count, subdivide it before scaling - if self.sub_idxs is not None and self.cond_hint_original.size(0) >= self.full_latent_length: - self.cond_hint = comfy.utils.common_upscale(self.cond_hint_original[self.sub_idxs], x_noisy.shape[3] * 8, x_noisy.shape[2] * 8, 'nearest-exact', "center").to(dtype).to(self.device) + if self.sub_idxs is not None: + actual_cond_hint_orig = self.cond_hint_original + if self.cond_hint_original.size(0) < self.full_latent_length: + actual_cond_hint_orig = extend_to_batch_size(tensor=actual_cond_hint_orig, batch_size=self.full_latent_length) + self.cond_hint = comfy.utils.common_upscale(actual_cond_hint_orig[self.sub_idxs], x_noisy.shape[3] * 8, x_noisy.shape[2] * 8, 'nearest-exact', "center").to(dtype).to(self.device) else: self.cond_hint = comfy.utils.common_upscale(self.cond_hint_original, x_noisy.shape[3] * 8, x_noisy.shape[2] * 8, 'nearest-exact', "center").to(dtype).to(self.device) 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) + self.cond_hint = broadcast_image_to_extend(self.cond_hint, x_noisy.shape[0], batched_number) # some special logic here compared to other controlnets: # * The cond_emb in attn patches will divide latent dims by 2 or 4, integer # * Due to this loss, the cond_emb will become smaller than x input if latent dims are not divisble by 2 or 4 diff --git a/adv_control/control_reference.py b/adv_control/control_reference.py index 815d651..24f1485 100644 --- a/adv_control/control_reference.py +++ b/adv_control/control_reference.py @@ -14,7 +14,7 @@ from comfy.ldm.modules.diffusionmodules import openaimodel from .logger import logger from .utils import (AdvancedControlBase, ControlWeights, TimestepKeyframeGroup, AbstractPreprocWrapper, - deepcopy_with_sharing, prepare_mask_batch, broadcast_image_to_full) + deepcopy_with_sharing, prepare_mask_batch, broadcast_image_to_extend) def refcn_sample_factory(orig_comfy_sample: Callable, is_custom=False) -> Callable: @@ -326,7 +326,7 @@ class ReferenceAdvanced(ControlBase, AdvancedControlBase): self.cond_hint_original, x_noisy.shape[3], x_noisy.shape[2], 'nearest-exact', "center").to(dtype).to(self.device) if x_noisy.shape[0] != self.cond_hint.shape[0]: - self.cond_hint = broadcast_image_to_full(self.cond_hint, x_noisy.shape[0], batched_number, except_one=False) + self.cond_hint = broadcast_image_to_extend(self.cond_hint, x_noisy.shape[0], batched_number, except_one=False) # noise cond_hint based on sigma (current step) self.cond_hint = self.latent_format.process_in(self.cond_hint) self.cond_hint = ref_noise_latents(self.cond_hint, sigma=t, noise=None) diff --git a/adv_control/control_sparsectrl.py b/adv_control/control_sparsectrl.py index 58b4d85..c50c7c7 100644 --- a/adv_control/control_sparsectrl.py +++ b/adv_control/control_sparsectrl.py @@ -25,12 +25,11 @@ from comfy.ldm.modules.attention import attention_basic, attention_pytorch, atte from comfy.ldm.modules.attention import FeedForward, SpatialTransformer from comfy.ldm.modules.diffusionmodules.openaimodel import TimestepEmbedSequential, ResBlock, Downsample from comfy.model_patcher import ModelPatcher -from comfy.controlnet import broadcast_image_to -from comfy.utils import repeat_to_batch_size import comfy.ops import comfy.model_management -from .utils import TimestepKeyframeGroup, disable_weight_init_clean_groupnorm, prepare_mask_batch +from .utils import (TimestepKeyframeGroup, disable_weight_init_clean_groupnorm, + prepare_mask_batch, broadcast_image_to_extend, extend_to_batch_size) # until xformers bug is fixed, do not use xformers for VersatileAttention! TODO: change this when fix is out @@ -560,8 +559,8 @@ class VanillaTemporalModule(nn.Module): return self.temporal_transformer(input_tensor, encoder_hidden_states, attention_mask) elif math.isclose(self.strength, 0.0): return input_tensor - elif self.strength > 1.0: - return self.temporal_transformer(input_tensor, encoder_hidden_states, attention_mask)*self.strength + # elif self.strength > 1.0: + # return self.temporal_transformer(input_tensor, encoder_hidden_states, attention_mask)*self.strength else: return self.temporal_transformer(input_tensor, encoder_hidden_states, attention_mask)*self.strength + input_tensor*(1.0-self.strength) @@ -673,10 +672,10 @@ class TemporalTransformer3DModel(nn.Module): # otherwise, calculate temp mask self.prev_hidden_states_batch = batch mask = prepare_mask_batch(self.raw_scale_mask, shape=(self.full_length, 1, height, width)) - mask = repeat_to_batch_size(mask, self.full_length) + mask = extend_to_batch_size(mask, self.full_length) # if mask not the same amount length as full length, make it match if self.full_length != mask.shape[0]: - mask = broadcast_image_to(mask, self.full_length, 1) + mask = broadcast_image_to_extend(mask, self.full_length, 1) # reshape mask to attention K shape (h*w, latent_count, 1) batch, channel, height, width = mask.shape # first, perform same operations as on hidden_states, diff --git a/adv_control/utils.py b/adv_control/utils.py index 9e0beda..a7e17dd 100644 --- a/adv_control/utils.py +++ b/adv_control/utils.py @@ -12,7 +12,7 @@ import comfy.sample import comfy.samplers import comfy.model_base -from comfy.controlnet import ControlBase, broadcast_image_to +from comfy.controlnet import ControlBase from comfy.model_patcher import ModelPatcher from .logger import logger @@ -152,7 +152,7 @@ class ControlWeightType: class ControlWeights: def __init__(self, weight_type: str, base_multiplier: float=1.0, flip_weights: bool=False, weights: list[float]=None, weight_mask: Tensor=None, - uncond_multiplier=1.0): + uncond_multiplier=1.0, uncond_mask: Tensor=None): self.weight_type = weight_type self.base_multiplier = base_multiplier self.flip_weights = flip_weights @@ -162,6 +162,8 @@ class ControlWeights: self.weight_mask = weight_mask self.uncond_multiplier = float(uncond_multiplier) self.has_uncond_multiplier = not math.isclose(self.uncond_multiplier, 1.0) + self.uncond_mask = uncond_mask if uncond_mask is not None else 1.0 + self.has_uncond_mask = uncond_mask is not None def get(self, idx: int, default=1.0) -> Union[float, Tensor]: # if weights is not none, return index @@ -433,8 +435,15 @@ def normalize_min_max(x: Tensor, new_min = 0.0, new_max = 1.0): def linear_conversion(x, x_min=0.0, x_max=1.0, new_min=0.0, new_max=1.0): return (((x - x_min)/(x_max - x_min)) * (new_max - new_min)) + new_min +def extend_to_batch_size(tensor: Tensor, batch_size: int): + if tensor.shape[0] > batch_size: + return tensor[:batch_size] + elif tensor.shape[0] < batch_size: + remainder = batch_size-tensor.shape[0] + return torch.cat([tensor] + [tensor[-1:]]*remainder, dim=0) + return tensor -def broadcast_image_to_full(tensor, target_batch_size, batched_number, except_one=True): +def broadcast_image_to_extend(tensor, target_batch_size, batched_number, except_one=True): current_batch_size = tensor.shape[0] #print(current_batch_size, target_batch_size) if except_one and current_batch_size == 1: @@ -444,7 +453,7 @@ def broadcast_image_to_full(tensor, target_batch_size, batched_number, except_on tensor = tensor[:per_batch] if per_batch > tensor.shape[0]: - tensor = torch.cat([tensor] * (per_batch // tensor.shape[0]) + [tensor[:(per_batch % tensor.shape[0])]], dim=0) + tensor = extend_to_batch_size(tensor=tensor, batch_size=per_batch) current_batch_size = tensor.shape[0] if current_batch_size == target_batch_size: @@ -772,6 +781,8 @@ class AdvancedControlBase: # if uncond, set to weight's uncond_multiplier if cond_type == 1: x[actual_length*idx:actual_length*(idx+1)] *= self.weights.uncond_multiplier + if self.weights.has_uncond_mask: + pass if self.latent_keyframes is not None: x[:] = x[:] * self.calc_latent_keyframe_mults(x=x, batched_number=batched_number) @@ -860,12 +871,12 @@ class AdvancedControlBase: # resize mask and match batch count out_mask = prepare_mask_batch(orig_mask, x_noisy.shape, multiplier=multiplier) actual_latent_length = x_noisy.shape[0] // batched_number - out_mask = comfy.utils.repeat_to_batch_size(out_mask, actual_latent_length if self.sub_idxs is None else self.full_latent_length) + out_mask = extend_to_batch_size(out_mask, actual_latent_length if self.sub_idxs is None else self.full_latent_length) if self.sub_idxs is not None: out_mask = out_mask[self.sub_idxs] # make cond_hint_mask length match x_noise if x_noisy.shape[0] != out_mask.shape[0]: - out_mask = broadcast_image_to(out_mask, x_noisy.shape[0], batched_number) + out_mask = broadcast_image_to_extend(out_mask, x_noisy.shape[0], batched_number) # default dtype to be same as x_noisy if dtype is None: dtype = x_noisy.dtype From 20d54c06608727012babcd4a83b54ca28656553d Mon Sep 17 00:00:00 2001 From: Jedrzej Kosinski Date: Mon, 27 May 2024 18:49:28 -0500 Subject: [PATCH 4/9] Added support for weight extras, added sparse_hint_mult, sparse_nonhint_mult, and sparse_mask_strength for SparseCtrl, renamed Force Default Weights node to Default Weights --- adv_control/control.py | 24 ++++++++++++++++--- adv_control/control_sparsectrl.py | 13 ++++++++++- adv_control/nodes.py | 6 +++-- adv_control/nodes_sparsectrl.py | 39 ++++++++++++++++++++++++++++--- adv_control/nodes_weight.py | 37 ++++++++++++++++++----------- adv_control/utils.py | 39 ++++++++++++++++--------------- 6 files changed, 116 insertions(+), 42 deletions(-) diff --git a/adv_control/control.py b/adv_control/control.py index 235bb02..cab2470 100644 --- a/adv_control/control.py +++ b/adv_control/control.py @@ -11,7 +11,7 @@ import comfy.controlnet as comfy_cn from comfy.controlnet import ControlBase, ControlNet, ControlLora, T2IAdapter from comfy.model_patcher import ModelPatcher -from .control_sparsectrl import SparseModelPatcher, SparseControlNet, SparseCtrlMotionWrapper, SparseMethod, SparseSettings, SparseSpreadMethod, PreprocSparseRGBWrapper +from .control_sparsectrl import SparseModelPatcher, SparseControlNet, SparseCtrlMotionWrapper, SparseMethod, SparseSettings, SparseSpreadMethod, PreprocSparseRGBWrapper, SparseConst from .control_lllite import LLLiteModule, LLLitePatch from .control_svd import svd_unet_config_from_diffusers_unet, SVDControlNet, svd_unet_to_diffusers from .utils import (AdvancedControlBase, TimestepKeyframeGroup, LatentKeyframeGroup, ControlWeightType, ControlWeights, WeightTypeException, @@ -306,11 +306,20 @@ class SparseCtrlAdvanced(ControlNetAdvanced): range_idxs = list(range(full_length)) if self.sub_idxs is None else self.sub_idxs hint_idxs = [] # idxs in cond_idxs - local_idxs = [] # idx to pun in final cond_hint + local_idxs = [] # idx to put in final cond_hint for i,cond_idx in enumerate(cond_idxs): if cond_idx in range_idxs: hint_idxs.append(i) local_idxs.append(range_idxs.index(cond_idx)) + # determine cond/uncond indexes that will get masked + self.local_sparse_idxs = [] + self.local_sparse_idxs_inverse = list(range(x_noisy.size(0))) + for batch_idx in range(batched_number): + for i in local_idxs: + actual_i = i+(batch_idx*actual_length) + self.local_sparse_idxs.append(actual_i) + if actual_i in self.local_sparse_idxs_inverse: + self.local_sparse_idxs_inverse.remove(actual_i) # sub_cond_hint now contains the hints relevant to current x_noisy sub_cond_hint = self.cond_hint_original[hint_idxs].to(dtype).to(self.device) @@ -330,7 +339,8 @@ class SparseCtrlAdvanced(ControlNetAdvanced): # prepare cond_mask (b, 1, h, w) cond_shape[1] = 1 cond_mask = torch.zeros(cond_shape).to(dtype).to(self.device) - cond_mask[local_idxs] = 1.0 + cond_mask[local_idxs] = self.sparse_settings.sparse_mask_strength * self.weights.extras.get(SparseConst.MASK_STRENGTH, 1.0) + #cond_mask[local_idxs] = 2.5 # combine cond_hint and cond_mask into (b, c+1, h, w) if not self.sparse_settings.merged: self.cond_hint = torch.cat([self.cond_hint, cond_mask], dim=1) @@ -353,6 +363,12 @@ class SparseCtrlAdvanced(ControlNetAdvanced): control = self.control_model(x=x_noisy.to(dtype), hint=self.cond_hint, timesteps=timestep.float(), context=context.to(dtype), y=y) return self.control_merge(None, control, control_prev, output_dtype) + def apply_advanced_strengths_and_masks(self, x: Tensor, batched_number: int): + # apply mults to indexes with and without a direct condhint + x[self.local_sparse_idxs] *= self.sparse_settings.sparse_hint_mult * self.weights.extras.get(SparseConst.HINT_MULT, 1.0) + x[self.local_sparse_idxs_inverse] *= self.sparse_settings.sparse_nonhint_mult * self.weights.extras.get(SparseConst.NONHINT_MULT, 1.0) + return super().apply_advanced_strengths_and_masks(x, batched_number) + def pre_run_advanced(self, model, percent_to_timestep_function): super().pre_run_advanced(model, percent_to_timestep_function) if type(self.cond_hint_original) == PreprocSparseRGBWrapper: @@ -370,6 +386,8 @@ class SparseCtrlAdvanced(ControlNetAdvanced): if self.latent_format is not None: del self.latent_format self.latent_format = None + self.local_sparse_idxs = None + self.local_sparse_idxs_inverse = None def copy(self): c = SparseCtrlAdvanced(self.control_model, self.timestep_keyframes, self.sparse_settings, self.global_average_pooling, self.device, self.load_device, self.manual_cast_dtype) diff --git a/adv_control/control_sparsectrl.py b/adv_control/control_sparsectrl.py index c50c7c7..dbc2112 100644 --- a/adv_control/control_sparsectrl.py +++ b/adv_control/control_sparsectrl.py @@ -53,6 +53,12 @@ else: optimized_attention_mm = attention_sub_quad +class SparseConst: + HINT_MULT = "sparse_hint_mult" + NONHINT_MULT = "sparse_nonhint_mult" + MASK_STRENGTH = "sparse_mask_strength" + + class SparseControlNet(ControlNetCLDM): def __init__(self, *args,**kwargs): super().__init__(*args, **kwargs) @@ -179,12 +185,17 @@ class PreprocSparseRGBWrapper: class SparseSettings: - def __init__(self, sparse_method: 'SparseMethod', use_motion: bool=True, motion_strength=1.0, motion_scale=1.0, merged=False): + def __init__(self, sparse_method: 'SparseMethod', use_motion: bool=True, motion_strength=1.0, motion_scale=1.0, merged=False, + sparse_mask_strength=1.0, sparse_hint_mult=1.0, sparse_nonhint_mult=1.0, context_aware=True): self.sparse_method = sparse_method self.use_motion = use_motion self.motion_strength = motion_strength self.motion_scale = motion_scale self.merged = merged + self.sparse_mask_strength = float(sparse_mask_strength) + self.sparse_hint_mult = float(sparse_hint_mult) + self.sparse_nonhint_mult = float(sparse_nonhint_mult) + self.context_aware = context_aware @classmethod def default(cls): diff --git a/adv_control/nodes.py b/adv_control/nodes.py index 5d56554..14659c7 100644 --- a/adv_control/nodes.py +++ b/adv_control/nodes.py @@ -10,7 +10,7 @@ from .nodes_weight import (DefaultWeights, ScaledSoftMaskedUniversalWeights, Sca SoftT2IAdapterWeights, CustomT2IAdapterWeights) from .nodes_keyframes import (LatentKeyframeGroupNode, LatentKeyframeInterpolationNode, LatentKeyframeBatchedGroupNode, LatentKeyframeNode, TimestepKeyframeNode, TimestepKeyframeInterpolationNode, TimestepKeyframeFromStrengthListNode) -from .nodes_sparsectrl import SparseCtrlMergedLoaderAdvanced, SparseCtrlLoaderAdvanced, SparseIndexMethodNode, SparseSpreadMethodNode, RgbSparseCtrlPreprocessor +from .nodes_sparsectrl import SparseCtrlMergedLoaderAdvanced, SparseCtrlLoaderAdvanced, SparseIndexMethodNode, SparseSpreadMethodNode, RgbSparseCtrlPreprocessor, SparseWeightExtras from .nodes_reference import ReferenceControlNetNode, ReferenceControlFinetune, ReferencePreprocessorNode from .nodes_loosecontrol import ControlNetLoaderWithLoraAdvanced from .nodes_deprecated import LoadImagesFromDirectory @@ -186,6 +186,7 @@ NODE_CLASS_MAPPINGS = { "ACN_SparseCtrlMergedLoaderAdvanced": SparseCtrlMergedLoaderAdvanced, "ACN_SparseCtrlIndexMethodNode": SparseIndexMethodNode, "ACN_SparseCtrlSpreadMethodNode": SparseSpreadMethodNode, + "ACN_SparseCtrlWeightExtras": SparseWeightExtras, # Reference "ACN_ReferencePreprocessor": ReferencePreprocessorNode, "ACN_ReferenceControlNet": ReferenceControlNetNode, @@ -217,13 +218,14 @@ NODE_DISPLAY_NAME_MAPPINGS = { "CustomControlNetWeights": "ControlNet Custom Weights ๐Ÿ›‚๐Ÿ…๐Ÿ…’๐Ÿ…", "SoftT2IAdapterWeights": "T2IAdapter Soft Weights ๐Ÿ›‚๐Ÿ…๐Ÿ…’๐Ÿ…", "CustomT2IAdapterWeights": "T2IAdapter Custom Weights ๐Ÿ›‚๐Ÿ…๐Ÿ…’๐Ÿ…", - "ACN_DefaultUniversalWeights": "Force Default Weights ๐Ÿ›‚๐Ÿ…๐Ÿ…’๐Ÿ…", + "ACN_DefaultUniversalWeights": "Default Weights ๐Ÿ›‚๐Ÿ…๐Ÿ…’๐Ÿ…", # SparseCtrl "ACN_SparseCtrlRGBPreprocessor": "RGB SparseCtrl ๐Ÿ›‚๐Ÿ…๐Ÿ…’๐Ÿ…", "ACN_SparseCtrlLoaderAdvanced": "Load SparseCtrl Model ๐Ÿ›‚๐Ÿ…๐Ÿ…’๐Ÿ…", "ACN_SparseCtrlMergedLoaderAdvanced": "๐ŸงชLoad Merged SparseCtrl Model ๐Ÿ›‚๐Ÿ…๐Ÿ…’๐Ÿ…", "ACN_SparseCtrlIndexMethodNode": "SparseCtrl Index Method ๐Ÿ›‚๐Ÿ…๐Ÿ…’๐Ÿ…", "ACN_SparseCtrlSpreadMethodNode": "SparseCtrl Spread Method ๐Ÿ›‚๐Ÿ…๐Ÿ…’๐Ÿ…", + "ACN_SparseCtrlWeightExtras": "SparseCtrl Weight Extras ๐Ÿ›‚๐Ÿ…๐Ÿ…’๐Ÿ…", # Reference "ACN_ReferencePreprocessor": "Reference Preproccessor ๐Ÿ›‚๐Ÿ…๐Ÿ…’๐Ÿ…", "ACN_ReferenceControlNet": "Reference ControlNet ๐Ÿ›‚๐Ÿ…๐Ÿ…’๐Ÿ…", diff --git a/adv_control/nodes_sparsectrl.py b/adv_control/nodes_sparsectrl.py index 4df32b0..cfbfd2a 100644 --- a/adv_control/nodes_sparsectrl.py +++ b/adv_control/nodes_sparsectrl.py @@ -6,7 +6,7 @@ import comfy.utils from comfy.sd import VAE from .utils import TimestepKeyframeGroup -from .control_sparsectrl import SparseMethod, SparseIndexMethod, SparseSettings, SparseSpreadMethod, PreprocSparseRGBWrapper +from .control_sparsectrl import SparseMethod, SparseIndexMethod, SparseSettings, SparseSpreadMethod, PreprocSparseRGBWrapper, SparseConst from .control import load_sparsectrl, load_controlnet, ControlNetAdvanced, SparseCtrlAdvanced @@ -24,6 +24,10 @@ class SparseCtrlLoaderAdvanced: "optional": { "sparse_method": ("SPARSE_METHOD", ), "tk_optional": ("TIMESTEP_KEYFRAME", ), + "context_aware": ("BOOLEAN", {"default": True}, ), + "sparse_hint_mult": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 10.0, "step": 0.001}, ), + "sparse_nonhint_mult": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 10.0, "step": 0.001}, ), + "sparse_mask_strength": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 10.0, "step": 0.001}, ), } } @@ -32,9 +36,12 @@ class SparseCtrlLoaderAdvanced: CATEGORY = "Adv-ControlNet ๐Ÿ›‚๐Ÿ…๐Ÿ…’๐Ÿ…/SparseCtrl" - def load_controlnet(self, sparsectrl_name: str, use_motion: bool, motion_strength: float, motion_scale: float, sparse_method: SparseMethod=SparseSpreadMethod(), tk_optional: TimestepKeyframeGroup=None): + def load_controlnet(self, sparsectrl_name: str, use_motion: bool, motion_strength: float, motion_scale: float, sparse_method: SparseMethod=SparseSpreadMethod(), tk_optional: TimestepKeyframeGroup=None, + context_aware=True, sparse_hint_mult=1.0, sparse_nonhint_mult=1.0, sparse_mask_strength=1.0): sparsectrl_path = folder_paths.get_full_path("controlnet", sparsectrl_name) - sparse_settings = SparseSettings(sparse_method=sparse_method, use_motion=use_motion, motion_strength=motion_strength, motion_scale=motion_scale) + sparse_settings = SparseSettings(sparse_method=sparse_method, use_motion=use_motion, motion_strength=motion_strength, motion_scale=motion_scale, + context_aware=context_aware, + sparse_mask_strength=sparse_mask_strength, sparse_hint_mult=sparse_hint_mult, sparse_nonhint_mult=sparse_nonhint_mult) sparsectrl = load_sparsectrl(sparsectrl_path, timestep_keyframe=tk_optional, sparse_settings=sparse_settings) return (sparsectrl,) @@ -161,3 +168,29 @@ class RgbSparseCtrlPreprocessor: image = VAEEncode.vae_encode_crop_pixels(image) encoded = vae.encode(image[:,:,:,:3]) return (PreprocSparseRGBWrapper(condhint=encoded),) + + +class SparseWeightExtras: + @classmethod + def INPUT_TYPES(s): + return { + "optional": { + "extras": ("CN_WEIGHTS_EXTRAS",), + "sparse_hint_mult": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 10.0, "step": 0.001}, ), + "sparse_nonhint_mult": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 10.0, "step": 0.001}, ), + "sparse_mask_strength": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 10.0, "step": 0.001}, ), + } + } + + RETURN_TYPES = ("CN_WEIGHTS_EXTRAS", ) + RETURN_NAMES = ("extras", ) + FUNCTION = "create_weight_extras" + + CATEGORY = "Adv-ControlNet ๐Ÿ›‚๐Ÿ…๐Ÿ…’๐Ÿ…/SparseCtrl" + + def create_weight_extras(self, extras: dict[str]={}, sparse_hint_mult=1.0, sparse_nonhint_mult=1.0, sparse_mask_strength=1.0): + extras = extras.copy() + extras[SparseConst.HINT_MULT] = sparse_hint_mult + extras[SparseConst.NONHINT_MULT] = sparse_nonhint_mult + extras[SparseConst.MASK_STRENGTH] = sparse_mask_strength + return (extras, ) diff --git a/adv_control/nodes_weight.py b/adv_control/nodes_weight.py index 7eb58d4..ddd7ef0 100644 --- a/adv_control/nodes_weight.py +++ b/adv_control/nodes_weight.py @@ -11,6 +11,9 @@ class DefaultWeights: @classmethod def INPUT_TYPES(s): return { + "optional": { + "extras": ("CN_WEIGHTS_EXTRAS",), + } } RETURN_TYPES = ("CONTROL_NET_WEIGHTS", "TIMESTEP_KEYFRAME",) @@ -19,8 +22,8 @@ class DefaultWeights: CATEGORY = "Adv-ControlNet ๐Ÿ›‚๐Ÿ…๐Ÿ…’๐Ÿ…/weights" - def load_weights(self): - weights = ControlWeights.default() + def load_weights(self, extras: dict[str]={}): + weights = ControlWeights.default(extras=extras) return (weights, TimestepKeyframeGroup.default(TimestepKeyframe(control_weights=weights))) @@ -37,6 +40,7 @@ class ScaledSoftMaskedUniversalWeights: }, "optional": { "uncond_multiplier": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01}, ), + "extras": ("CN_WEIGHTS_EXTRAS",), } } @@ -47,7 +51,7 @@ class ScaledSoftMaskedUniversalWeights: CATEGORY = "Adv-ControlNet ๐Ÿ›‚๐Ÿ…๐Ÿ…’๐Ÿ…/weights" def load_weights(self, mask: Tensor, min_base_multiplier: float, max_base_multiplier: float, lock_min=False, lock_max=False, - uncond_multiplier: float=1.0): + uncond_multiplier: float=1.0, extras: dict[str]={}): # normalize mask mask = mask.clone() x_min = 0.0 if lock_min else mask.min() @@ -56,7 +60,7 @@ class ScaledSoftMaskedUniversalWeights: mask = torch.ones_like(mask) * max_base_multiplier else: mask = linear_conversion(mask, x_min, x_max, min_base_multiplier, max_base_multiplier) - weights = ControlWeights.universal_mask(weight_mask=mask, uncond_multiplier=uncond_multiplier) + weights = ControlWeights.universal_mask(weight_mask=mask, uncond_multiplier=uncond_multiplier, extras=extras) return (weights, TimestepKeyframeGroup.default(TimestepKeyframe(control_weights=weights))) @@ -70,6 +74,7 @@ class ScaledSoftUniversalWeights: }, "optional": { "uncond_multiplier": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01}, ), + "extras": ("CN_WEIGHTS_EXTRAS",), } } @@ -79,8 +84,8 @@ class ScaledSoftUniversalWeights: CATEGORY = "Adv-ControlNet ๐Ÿ›‚๐Ÿ…๐Ÿ…’๐Ÿ…/weights" - def load_weights(self, base_multiplier, flip_weights, uncond_multiplier: float=1.0): - weights = ControlWeights.universal(base_multiplier=base_multiplier, flip_weights=flip_weights, uncond_multiplier=uncond_multiplier) + def load_weights(self, base_multiplier, flip_weights, uncond_multiplier: float=1.0, extras: dict[str]={}): + weights = ControlWeights.universal(base_multiplier=base_multiplier, flip_weights=flip_weights, uncond_multiplier=uncond_multiplier, extras=extras) return (weights, TimestepKeyframeGroup.default(TimestepKeyframe(control_weights=weights))) @@ -106,6 +111,7 @@ class SoftControlNetWeights: }, "optional": { "uncond_multiplier": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01}, ), + "extras": ("CN_WEIGHTS_EXTRAS",), } } @@ -117,10 +123,10 @@ class SoftControlNetWeights: def load_weights(self, weight_00, weight_01, weight_02, weight_03, weight_04, weight_05, weight_06, weight_07, weight_08, weight_09, weight_10, weight_11, weight_12, flip_weights, - uncond_multiplier: float=1.0): + uncond_multiplier: float=1.0, extras: dict[str]={}): weights = [weight_00, weight_01, weight_02, weight_03, weight_04, weight_05, weight_06, weight_07, weight_08, weight_09, weight_10, weight_11, weight_12] - weights = ControlWeights.controlnet(weights, flip_weights=flip_weights, uncond_multiplier=uncond_multiplier) + weights = ControlWeights.controlnet(weights, flip_weights=flip_weights, uncond_multiplier=uncond_multiplier, extras=extras) return (weights, TimestepKeyframeGroup.default(TimestepKeyframe(control_weights=weights))) @@ -146,6 +152,7 @@ class CustomControlNetWeights: }, "optional": { "uncond_multiplier": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01}, ), + "extras": ("CN_WEIGHTS_EXTRAS",), } } @@ -157,10 +164,10 @@ class CustomControlNetWeights: def load_weights(self, weight_00, weight_01, weight_02, weight_03, weight_04, weight_05, weight_06, weight_07, weight_08, weight_09, weight_10, weight_11, weight_12, flip_weights, - uncond_multiplier: float=1.0): + uncond_multiplier: float=1.0, extras: dict[str]={}): weights = [weight_00, weight_01, weight_02, weight_03, weight_04, weight_05, weight_06, weight_07, weight_08, weight_09, weight_10, weight_11, weight_12] - weights = ControlWeights.controlnet(weights, flip_weights=flip_weights, uncond_multiplier=uncond_multiplier) + weights = ControlWeights.controlnet(weights, flip_weights=flip_weights, uncond_multiplier=uncond_multiplier, extras=extras) return (weights, TimestepKeyframeGroup.default(TimestepKeyframe(control_weights=weights))) @@ -177,6 +184,7 @@ class SoftT2IAdapterWeights: }, "optional": { "uncond_multiplier": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01}, ), + "extras": ("CN_WEIGHTS_EXTRAS",), } } @@ -187,10 +195,10 @@ class SoftT2IAdapterWeights: CATEGORY = "Adv-ControlNet ๐Ÿ›‚๐Ÿ…๐Ÿ…’๐Ÿ…/weights/T2IAdapter" def load_weights(self, weight_00, weight_01, weight_02, weight_03, flip_weights, - uncond_multiplier: float=1.0): + uncond_multiplier: float=1.0, extras: dict[str]={}): weights = [weight_00, weight_01, weight_02, weight_03] weights = get_properly_arranged_t2i_weights(weights) - weights = ControlWeights.t2iadapter(weights, flip_weights=flip_weights, uncond_multiplier=uncond_multiplier) + weights = ControlWeights.t2iadapter(weights, flip_weights=flip_weights, uncond_multiplier=uncond_multiplier, extras=extras) return (weights, TimestepKeyframeGroup.default(TimestepKeyframe(control_weights=weights))) @@ -207,6 +215,7 @@ class CustomT2IAdapterWeights: }, "optional": { "uncond_multiplier": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01}, ), + "extras": ("CN_WEIGHTS_EXTRAS",), } } @@ -217,8 +226,8 @@ class CustomT2IAdapterWeights: CATEGORY = "Adv-ControlNet ๐Ÿ›‚๐Ÿ…๐Ÿ…’๐Ÿ…/weights/T2IAdapter" def load_weights(self, weight_00, weight_01, weight_02, weight_03, flip_weights, - uncond_multiplier: float=1.0): + uncond_multiplier: float=1.0, extras: dict[str]={}): weights = [weight_00, weight_01, weight_02, weight_03] weights = get_properly_arranged_t2i_weights(weights) - weights = ControlWeights.t2iadapter(weights, flip_weights=flip_weights, uncond_multiplier=uncond_multiplier) + weights = ControlWeights.t2iadapter(weights, flip_weights=flip_weights, uncond_multiplier=uncond_multiplier, extras=extras) return (weights, TimestepKeyframeGroup.default(TimestepKeyframe(control_weights=weights))) diff --git a/adv_control/utils.py b/adv_control/utils.py index a7e17dd..a92998e 100644 --- a/adv_control/utils.py +++ b/adv_control/utils.py @@ -152,7 +152,7 @@ class ControlWeightType: class ControlWeights: def __init__(self, weight_type: str, base_multiplier: float=1.0, flip_weights: bool=False, weights: list[float]=None, weight_mask: Tensor=None, - uncond_multiplier=1.0, uncond_mask: Tensor=None): + uncond_multiplier=1.0, uncond_mask: Tensor=None, extras: dict[str]={},): self.weight_type = weight_type self.base_multiplier = base_multiplier self.flip_weights = flip_weights @@ -164,6 +164,7 @@ class ControlWeights: self.has_uncond_multiplier = not math.isclose(self.uncond_multiplier, 1.0) self.uncond_mask = uncond_mask if uncond_mask is not None else 1.0 self.has_uncond_mask = uncond_mask is not None + self.extras = extras def get(self, idx: int, default=1.0) -> Union[float, Tensor]: # if weights is not none, return index @@ -176,44 +177,44 @@ class ControlWeights: def copy_with_new_weights(self, new_weights: list[float]): return ControlWeights(weight_type=self.weight_type, base_multiplier=self.base_multiplier, flip_weights=self.flip_weights, - weights=new_weights, weight_mask=self.weight_mask, uncond_multiplier=self.uncond_multiplier) + weights=new_weights, weight_mask=self.weight_mask, uncond_multiplier=self.uncond_multiplier, extras=self.extras) @classmethod - def default(cls): - return cls(ControlWeightType.DEFAULT) + def default(cls, extras: dict[str]={}): + return cls(ControlWeightType.DEFAULT, extras=extras) @classmethod - def universal(cls, base_multiplier: float, flip_weights: bool=False, uncond_multiplier: float=1.0): - return cls(ControlWeightType.UNIVERSAL, base_multiplier=base_multiplier, flip_weights=flip_weights, uncond_multiplier=uncond_multiplier) + def universal(cls, base_multiplier: float, flip_weights: bool=False, uncond_multiplier: float=1.0, extras: dict[str]={}): + return cls(ControlWeightType.UNIVERSAL, base_multiplier=base_multiplier, flip_weights=flip_weights, uncond_multiplier=uncond_multiplier, extras=extras) @classmethod - def universal_mask(cls, weight_mask: Tensor, uncond_multiplier: float=1.0): - return cls(ControlWeightType.UNIVERSAL, weight_mask=weight_mask, uncond_multiplier=uncond_multiplier) + def universal_mask(cls, weight_mask: Tensor, uncond_multiplier: float=1.0, extras: dict[str]={}): + return cls(ControlWeightType.UNIVERSAL, weight_mask=weight_mask, uncond_multiplier=uncond_multiplier, extras=extras) @classmethod - def t2iadapter(cls, weights: list[float]=None, flip_weights: bool=False, uncond_multiplier: float=1.0): + def t2iadapter(cls, weights: list[float]=None, flip_weights: bool=False, uncond_multiplier: float=1.0, extras: dict[str]={}): if weights is None: weights = [1.0]*12 - return cls(ControlWeightType.T2IADAPTER, weights=weights,flip_weights=flip_weights, uncond_multiplier=uncond_multiplier) + return cls(ControlWeightType.T2IADAPTER, weights=weights,flip_weights=flip_weights, uncond_multiplier=uncond_multiplier, extras=extras) @classmethod - def controlnet(cls, weights: list[float]=None, flip_weights: bool=False, uncond_multiplier: float=1.0): + def controlnet(cls, weights: list[float]=None, flip_weights: bool=False, uncond_multiplier: float=1.0, extras: dict[str]={}): if weights is None: weights = [1.0]*13 - return cls(ControlWeightType.CONTROLNET, weights=weights, flip_weights=flip_weights, uncond_multiplier=uncond_multiplier) + return cls(ControlWeightType.CONTROLNET, weights=weights, flip_weights=flip_weights, uncond_multiplier=uncond_multiplier, extras=extras) @classmethod - def controllora(cls, weights: list[float]=None, flip_weights: bool=False, uncond_multiplier: float=1.0): + def controllora(cls, weights: list[float]=None, flip_weights: bool=False, uncond_multiplier: float=1.0, extras: dict[str]={}): if weights is None: weights = [1.0]*10 - return cls(ControlWeightType.CONTROLLORA, weights=weights, flip_weights=flip_weights, uncond_multiplier=uncond_multiplier) + return cls(ControlWeightType.CONTROLLORA, weights=weights, flip_weights=flip_weights, uncond_multiplier=uncond_multiplier, extras=extras) @classmethod - def controllllite(cls, weights: list[float]=None, flip_weights: bool=False, uncond_multiplier: float=1.0): + def controllllite(cls, weights: list[float]=None, flip_weights: bool=False, uncond_multiplier: float=1.0, extras: dict[str]={}): if weights is None: # TODO: make this have a real value weights = [1.0]*200 - return cls(ControlWeightType.CONTROLLLLITE, weights=weights, flip_weights=flip_weights, uncond_multiplier=uncond_multiplier) + return cls(ControlWeightType.CONTROLLLLITE, weights=weights, flip_weights=flip_weights, uncond_multiplier=uncond_multiplier, extras=extras) class StrengthInterpolation: @@ -532,7 +533,7 @@ class WeightTypeException(TypeError): class AdvancedControlBase: def __init__(self, base: ControlBase, timestep_keyframes: TimestepKeyframeGroup, weights_default: ControlWeights, require_model=False): self.base = base - self.compatible_weights = [ControlWeightType.UNIVERSAL] + self.compatible_weights = [ControlWeightType.UNIVERSAL, ControlWeightType.DEFAULT] self.add_compatible_weight(weights_default.weight_type) # mask for which parts of controlnet output to keep self.mask_cond_hint_original = None @@ -585,7 +586,7 @@ class AdvancedControlBase: else: for tk in self.timestep_keyframes.keyframes: if tk.has_control_weights() and tk.control_weights.weight_type not in self.compatible_weights: - msg = f"Weight on Timestep Keyframe with start_percent={tk.start_percent} is type" + \ + msg = f"Weight on Timestep Keyframe with start_percent={tk.start_percent} is type " + \ f"{tk.control_weights.weight_type}, but loaded {type(self).__name__} only supports {self.compatible_weights} weights." raise WeightTypeException(msg) @@ -652,7 +653,7 @@ class AdvancedControlBase: self.prepare_weights() def prepare_weights(self): - if self.weights is None or self.weights.weight_type == ControlWeightType.DEFAULT: + if self.weights is None: self.weights = self.weights_default elif self.weights.weight_type == ControlWeightType.UNIVERSAL: # if universal and weight_mask present, no need to convert From 74b6cf576fc8ebf8e7ac162cc1d573d4b93145c1 Mon Sep 17 00:00:00 2001 From: Jedrzej Kosinski Date: Tue, 28 May 2024 22:32:19 -0500 Subject: [PATCH 5/9] Implemented context_aware --- adv_control/control.py | 14 +++++-- adv_control/control_sparsectrl.py | 67 ++++++++++++++++++++++++++++--- 2 files changed, 71 insertions(+), 10 deletions(-) diff --git a/adv_control/control.py b/adv_control/control.py index cab2470..787277f 100644 --- a/adv_control/control.py +++ b/adv_control/control.py @@ -302,8 +302,8 @@ class SparseCtrlAdvanced(ControlNetAdvanced): del self.cond_hint self.cond_hint = None # first, figure out which cond idxs are relevant, and where they fit in - cond_idxs = self.sparse_settings.sparse_method.get_indexes(hint_length=self.cond_hint_original.size(0), full_length=full_length) - + cond_idxs, hint_order = self.sparse_settings.sparse_method.get_indexes(hint_length=self.cond_hint_original.size(0), full_length=full_length, + sub_idxs=self.sub_idxs if self.sparse_settings.context_aware else None) range_idxs = list(range(full_length)) if self.sub_idxs is None else self.sub_idxs hint_idxs = [] # idxs in cond_idxs local_idxs = [] # idx to put in final cond_hint @@ -311,6 +311,10 @@ class SparseCtrlAdvanced(ControlNetAdvanced): if cond_idx in range_idxs: hint_idxs.append(i) local_idxs.append(range_idxs.index(cond_idx)) + # log_string = f"cond_idxs: {cond_idxs}, local_idxs: {local_idxs}, hint_idxs: {hint_idxs}, hint_order: {hint_order}" + # if self.sub_idxs is not None: + # log_string += f" sub_idxs: {self.sub_idxs[0]}-{self.sub_idxs[-1]}" + # logger.warn(log_string) # determine cond/uncond indexes that will get masked self.local_sparse_idxs = [] self.local_sparse_idxs_inverse = list(range(x_noisy.size(0))) @@ -321,8 +325,10 @@ class SparseCtrlAdvanced(ControlNetAdvanced): if actual_i in self.local_sparse_idxs_inverse: self.local_sparse_idxs_inverse.remove(actual_i) # sub_cond_hint now contains the hints relevant to current x_noisy - sub_cond_hint = self.cond_hint_original[hint_idxs].to(dtype).to(self.device) - + if hint_order is None: + sub_cond_hint = self.cond_hint_original[hint_idxs].to(dtype).to(self.device) + else: + sub_cond_hint = self.cond_hint_original[hint_order][hint_idxs].to(dtype).to(self.device) # scale cond_hints to match noisy input if self.control_model.use_simplified_conditioning_embedding: # RGB SparseCtrl; the inputs are latents - use bilinear to avoid blocky artifacts diff --git a/adv_control/control_sparsectrl.py b/adv_control/control_sparsectrl.py index dbc2112..206b76d 100644 --- a/adv_control/control_sparsectrl.py +++ b/adv_control/control_sparsectrl.py @@ -28,7 +28,8 @@ from comfy.model_patcher import ModelPatcher import comfy.ops import comfy.model_management -from .utils import (TimestepKeyframeGroup, disable_weight_init_clean_groupnorm, +from .logger import logger +from .utils import (BIGMAX, TimestepKeyframeGroup, disable_weight_init_clean_groupnorm, prepare_mask_batch, broadcast_image_to_extend, extend_to_batch_size) @@ -209,7 +210,7 @@ class SparseMethod(ABC): self.method = method @abstractmethod - def get_indexes(self, hint_length: int, full_length: int) -> list[int]: + def get_indexes(self, hint_length: int, full_length: int, sub_idxs: list[int]=None) -> tuple[list[int], list[int]]: pass @@ -225,7 +226,60 @@ class SparseSpreadMethod(SparseMethod): super().__init__(self.SPREAD) self.spread = spread - def get_indexes(self, hint_length: int, full_length: int) -> list[int]: + def get_indexes(self, hint_length: int, full_length: int, sub_idxs: list[int]=None) -> tuple[list[int], list[int]]: + returned_idxs = self._get_indexes(hint_length, full_length) + if sub_idxs is None: + return returned_idxs, None + # need to map full indexes to condhint indexes + index_mapping = {} + for i, value in enumerate(returned_idxs): + index_mapping[value] = i + def get_mapped_idxs(idxs: list[int]): + return [index_mapping[idx] for idx in idxs] + # check if returned_idxs fit within subidxs + fitting_idxs = [] + for sub_idx in sub_idxs: + if sub_idx in returned_idxs: + fitting_idxs.append(sub_idx) + # if have any fitting_idxs, deal with it + if len(fitting_idxs) > 0: + return fitting_idxs, get_mapped_idxs(fitting_idxs) + + # since no returned_idxs fit in sub_idxs, need to get the next-closest hint images based on strategy + def get_closest_idx(target_idx: int, idxs: list[int]): + min_idx = -1 + min_dist = BIGMAX + for idx in idxs: + new_dist = abs(idx-target_idx) + if new_dist < min_dist: + min_idx = idx + min_dist = new_dist + if min_dist == 1: + return min_idx, min_dist + return min_idx, min_dist + start_closest_idx, start_dist = get_closest_idx(sub_idxs[0], returned_idxs) + end_closest_idx, end_dist = get_closest_idx(sub_idxs[-1], returned_idxs) + # if only one cond hint exists, do special behavior + if hint_length == 1: + # if same distance from start and end, + if start_dist == end_dist: + # find center index of sub_idxs + center_idx = sub_idxs[np.linspace(0, len(sub_idxs)-1, 3, endpoint=True, dtype=int)[1]] + return [center_idx], get_mapped_idxs([start_closest_idx]) + # otherwise, return closest + if start_dist < end_dist: + return [sub_idxs[0]], get_mapped_idxs([start_closest_idx]) + return [sub_idxs[-1]], get_mapped_idxs([end_closest_idx]) + # otherwise, select up to two closest images, or just 1, whichever one applies best + # if same distance from start and end, return two images to use + if start_dist == end_dist: + return [sub_idxs[0], sub_idxs[-1]], get_mapped_idxs([start_closest_idx, end_closest_idx]) + # else, use just one + if start_dist < end_dist: + return [sub_idxs[0]], get_mapped_idxs([start_closest_idx]) + return [sub_idxs[-1]], get_mapped_idxs([end_closest_idx]) + + def _get_indexes(self, hint_length: int, full_length: int) -> list[int]: # if hint_length >= full_length, limit hints to full_length if hint_length >= full_length: return list(range(full_length)) @@ -258,12 +312,13 @@ class SparseSpreadMethod(SparseMethod): return ValueError(f"Unrecognized spread: {self.spread}") -class SparseIndexMethod(SparseMethod): +class SparseIndexMethod(SparseSpreadMethod): def __init__(self, idxs: list[int]): - super().__init__(self.INDEX) + super().__init__() + self.method = self.INDEX self.idxs = idxs - def get_indexes(self, hint_length: int, full_length: int) -> list[int]: + def _get_indexes(self, hint_length: int, full_length: int) -> list[int]: orig_hint_length = hint_length if hint_length > full_length: hint_length = full_length From 5ef18991eeec36caa183d575c6c375088f538c2b Mon Sep 17 00:00:00 2001 From: Jedrzej Kosinski Date: Wed, 29 May 2024 01:00:54 -0500 Subject: [PATCH 6/9] Refactored SparseMethod code to have a common get_indexes func, with _get_indexes the one that is abstract --- adv_control/control_sparsectrl.py | 33 +++++++++++++++---------------- 1 file changed, 16 insertions(+), 17 deletions(-) diff --git a/adv_control/control_sparsectrl.py b/adv_control/control_sparsectrl.py index 206b76d..c70e3e2 100644 --- a/adv_control/control_sparsectrl.py +++ b/adv_control/control_sparsectrl.py @@ -210,22 +210,9 @@ class SparseMethod(ABC): self.method = method @abstractmethod - def get_indexes(self, hint_length: int, full_length: int, sub_idxs: list[int]=None) -> tuple[list[int], list[int]]: + def _get_indexes(self, hint_length: int, full_length: int) -> list[int]: pass - -class SparseSpreadMethod(SparseMethod): - UNIFORM = "uniform" - STARTING = "starting" - ENDING = "ending" - CENTER = "center" - - LIST = [UNIFORM, STARTING, ENDING, CENTER] - - def __init__(self, spread=UNIFORM): - super().__init__(self.SPREAD) - self.spread = spread - def get_indexes(self, hint_length: int, full_length: int, sub_idxs: list[int]=None) -> tuple[list[int], list[int]]: returned_idxs = self._get_indexes(hint_length, full_length) if sub_idxs is None: @@ -279,6 +266,19 @@ class SparseSpreadMethod(SparseMethod): return [sub_idxs[0]], get_mapped_idxs([start_closest_idx]) return [sub_idxs[-1]], get_mapped_idxs([end_closest_idx]) + +class SparseSpreadMethod(SparseMethod): + UNIFORM = "uniform" + STARTING = "starting" + ENDING = "ending" + CENTER = "center" + + LIST = [UNIFORM, STARTING, ENDING, CENTER] + + def __init__(self, spread=UNIFORM): + super().__init__(self.SPREAD) + self.spread = spread + def _get_indexes(self, hint_length: int, full_length: int) -> list[int]: # if hint_length >= full_length, limit hints to full_length if hint_length >= full_length: @@ -312,10 +312,9 @@ class SparseSpreadMethod(SparseMethod): return ValueError(f"Unrecognized spread: {self.spread}") -class SparseIndexMethod(SparseSpreadMethod): +class SparseIndexMethod(SparseMethod): def __init__(self, idxs: list[int]): - super().__init__() - self.method = self.INDEX + super().__init__(self.INDEX) self.idxs = idxs def _get_indexes(self, hint_length: int, full_length: int) -> list[int]: From 949e1acf7e0530a063370a3d493e6bf79d90c58c Mon Sep 17 00:00:00 2001 From: Jedrzej Kosinski Date: Wed, 29 May 2024 01:23:21 -0500 Subject: [PATCH 7/9] Turned context_aware into a dropdown in case I add other context_aware types in future, renamed sparse_mask_strength to sparse_mask_mult --- adv_control/control.py | 5 ++--- adv_control/control_sparsectrl.py | 16 +++++++++++++--- adv_control/nodes_sparsectrl.py | 18 +++++++++--------- 3 files changed, 24 insertions(+), 15 deletions(-) diff --git a/adv_control/control.py b/adv_control/control.py index 787277f..ce00c74 100644 --- a/adv_control/control.py +++ b/adv_control/control.py @@ -303,7 +303,7 @@ class SparseCtrlAdvanced(ControlNetAdvanced): self.cond_hint = None # first, figure out which cond idxs are relevant, and where they fit in cond_idxs, hint_order = self.sparse_settings.sparse_method.get_indexes(hint_length=self.cond_hint_original.size(0), full_length=full_length, - sub_idxs=self.sub_idxs if self.sparse_settings.context_aware else None) + sub_idxs=self.sub_idxs if self.sparse_settings.is_context_aware() else None) range_idxs = list(range(full_length)) if self.sub_idxs is None else self.sub_idxs hint_idxs = [] # idxs in cond_idxs local_idxs = [] # idx to put in final cond_hint @@ -345,8 +345,7 @@ class SparseCtrlAdvanced(ControlNetAdvanced): # prepare cond_mask (b, 1, h, w) cond_shape[1] = 1 cond_mask = torch.zeros(cond_shape).to(dtype).to(self.device) - cond_mask[local_idxs] = self.sparse_settings.sparse_mask_strength * self.weights.extras.get(SparseConst.MASK_STRENGTH, 1.0) - #cond_mask[local_idxs] = 2.5 + cond_mask[local_idxs] = self.sparse_settings.sparse_mask_mult * self.weights.extras.get(SparseConst.MASK_MULT, 1.0) # combine cond_hint and cond_mask into (b, c+1, h, w) if not self.sparse_settings.merged: self.cond_hint = torch.cat([self.cond_hint, cond_mask], dim=1) diff --git a/adv_control/control_sparsectrl.py b/adv_control/control_sparsectrl.py index c70e3e2..e10e08c 100644 --- a/adv_control/control_sparsectrl.py +++ b/adv_control/control_sparsectrl.py @@ -57,7 +57,7 @@ else: class SparseConst: HINT_MULT = "sparse_hint_mult" NONHINT_MULT = "sparse_nonhint_mult" - MASK_STRENGTH = "sparse_mask_strength" + MASK_MULT = "sparse_mask_mult" class SparseControlNet(ControlNetCLDM): @@ -185,19 +185,29 @@ class PreprocSparseRGBWrapper: raise AttributeError(self.error_msg) +class SparseContextAware: + NEAREST_HINT = "nearest_hint" + OFF = "off" + + LIST = [NEAREST_HINT, OFF] + + class SparseSettings: def __init__(self, sparse_method: 'SparseMethod', use_motion: bool=True, motion_strength=1.0, motion_scale=1.0, merged=False, - sparse_mask_strength=1.0, sparse_hint_mult=1.0, sparse_nonhint_mult=1.0, context_aware=True): + sparse_mask_mult=1.0, sparse_hint_mult=1.0, sparse_nonhint_mult=1.0, context_aware=SparseContextAware.NEAREST_HINT): self.sparse_method = sparse_method self.use_motion = use_motion self.motion_strength = motion_strength self.motion_scale = motion_scale self.merged = merged - self.sparse_mask_strength = float(sparse_mask_strength) + self.sparse_mask_mult = float(sparse_mask_mult) self.sparse_hint_mult = float(sparse_hint_mult) self.sparse_nonhint_mult = float(sparse_nonhint_mult) self.context_aware = context_aware + def is_context_aware(self): + return self.context_aware != SparseContextAware.OFF + @classmethod def default(cls): return SparseSettings(sparse_method=SparseSpreadMethod(), use_motion=True) diff --git a/adv_control/nodes_sparsectrl.py b/adv_control/nodes_sparsectrl.py index cfbfd2a..70eadc5 100644 --- a/adv_control/nodes_sparsectrl.py +++ b/adv_control/nodes_sparsectrl.py @@ -6,7 +6,7 @@ import comfy.utils from comfy.sd import VAE from .utils import TimestepKeyframeGroup -from .control_sparsectrl import SparseMethod, SparseIndexMethod, SparseSettings, SparseSpreadMethod, PreprocSparseRGBWrapper, SparseConst +from .control_sparsectrl import SparseMethod, SparseIndexMethod, SparseSettings, SparseSpreadMethod, PreprocSparseRGBWrapper, SparseConst, SparseContextAware from .control import load_sparsectrl, load_controlnet, ControlNetAdvanced, SparseCtrlAdvanced @@ -24,10 +24,10 @@ class SparseCtrlLoaderAdvanced: "optional": { "sparse_method": ("SPARSE_METHOD", ), "tk_optional": ("TIMESTEP_KEYFRAME", ), - "context_aware": ("BOOLEAN", {"default": True}, ), + "context_aware": (SparseContextAware.LIST, ), "sparse_hint_mult": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 10.0, "step": 0.001}, ), "sparse_nonhint_mult": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 10.0, "step": 0.001}, ), - "sparse_mask_strength": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 10.0, "step": 0.001}, ), + "sparse_mask_mult": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 10.0, "step": 0.001}, ), } } @@ -37,11 +37,11 @@ class SparseCtrlLoaderAdvanced: CATEGORY = "Adv-ControlNet ๐Ÿ›‚๐Ÿ…๐Ÿ…’๐Ÿ…/SparseCtrl" def load_controlnet(self, sparsectrl_name: str, use_motion: bool, motion_strength: float, motion_scale: float, sparse_method: SparseMethod=SparseSpreadMethod(), tk_optional: TimestepKeyframeGroup=None, - context_aware=True, sparse_hint_mult=1.0, sparse_nonhint_mult=1.0, sparse_mask_strength=1.0): + context_aware=SparseContextAware.NEAREST_HINT, sparse_hint_mult=1.0, sparse_nonhint_mult=1.0, sparse_mask_mult=1.0): sparsectrl_path = folder_paths.get_full_path("controlnet", sparsectrl_name) sparse_settings = SparseSettings(sparse_method=sparse_method, use_motion=use_motion, motion_strength=motion_strength, motion_scale=motion_scale, context_aware=context_aware, - sparse_mask_strength=sparse_mask_strength, sparse_hint_mult=sparse_hint_mult, sparse_nonhint_mult=sparse_nonhint_mult) + sparse_mask_mult=sparse_mask_mult, sparse_hint_mult=sparse_hint_mult, sparse_nonhint_mult=sparse_nonhint_mult) sparsectrl = load_sparsectrl(sparsectrl_path, timestep_keyframe=tk_optional, sparse_settings=sparse_settings) return (sparsectrl,) @@ -178,7 +178,7 @@ class SparseWeightExtras: "extras": ("CN_WEIGHTS_EXTRAS",), "sparse_hint_mult": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 10.0, "step": 0.001}, ), "sparse_nonhint_mult": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 10.0, "step": 0.001}, ), - "sparse_mask_strength": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 10.0, "step": 0.001}, ), + "sparse_mask_mult": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 10.0, "step": 0.001}, ), } } @@ -186,11 +186,11 @@ class SparseWeightExtras: RETURN_NAMES = ("extras", ) FUNCTION = "create_weight_extras" - CATEGORY = "Adv-ControlNet ๐Ÿ›‚๐Ÿ…๐Ÿ…’๐Ÿ…/SparseCtrl" + CATEGORY = "Adv-ControlNet ๐Ÿ›‚๐Ÿ…๐Ÿ…’๐Ÿ…/SparseCtrl/extras" - def create_weight_extras(self, extras: dict[str]={}, sparse_hint_mult=1.0, sparse_nonhint_mult=1.0, sparse_mask_strength=1.0): + def create_weight_extras(self, extras: dict[str]={}, sparse_hint_mult=1.0, sparse_nonhint_mult=1.0, sparse_mask_mult=1.0): extras = extras.copy() extras[SparseConst.HINT_MULT] = sparse_hint_mult extras[SparseConst.NONHINT_MULT] = sparse_nonhint_mult - extras[SparseConst.MASK_STRENGTH] = sparse_mask_strength + extras[SparseConst.MASK_MULT] = sparse_mask_mult return (extras, ) From 6940e3fd4b36a0eaa524e98c358ec23c1442ab2a Mon Sep 17 00:00:00 2001 From: Jedrzej Kosinski Date: Wed, 29 May 2024 01:23:44 -0500 Subject: [PATCH 8/9] version bump --- pyproject.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pyproject.toml b/pyproject.toml index 9df3e69..d3d8082 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,7 +1,7 @@ [project] name = "comfyui-advanced-controlnet" description = "Nodes for scheduling ControlNet strength across timesteps and batched latents, as well as applying custom weights and attention masks." -version = "1.0.2" +version = "1.0.3" license = "LICENSE" dependencies = [] From 576426a73495609f7a565db8dd25d77c670574a3 Mon Sep 17 00:00:00 2001 From: Jedrzej Kosinski Date: Wed, 29 May 2024 01:37:56 -0500 Subject: [PATCH 9/9] Renamed extras to cn_extras to differentiate from potential extras on AnimateDiff in the future (which would be called ad_extras there) --- adv_control/nodes_sparsectrl.py | 16 ++++++------- adv_control/nodes_weight.py | 42 ++++++++++++++++----------------- 2 files changed, 29 insertions(+), 29 deletions(-) diff --git a/adv_control/nodes_sparsectrl.py b/adv_control/nodes_sparsectrl.py index 70eadc5..9fde905 100644 --- a/adv_control/nodes_sparsectrl.py +++ b/adv_control/nodes_sparsectrl.py @@ -175,7 +175,7 @@ class SparseWeightExtras: def INPUT_TYPES(s): return { "optional": { - "extras": ("CN_WEIGHTS_EXTRAS",), + "cn_extras": ("CN_WEIGHTS_EXTRAS",), "sparse_hint_mult": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 10.0, "step": 0.001}, ), "sparse_nonhint_mult": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 10.0, "step": 0.001}, ), "sparse_mask_mult": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 10.0, "step": 0.001}, ), @@ -183,14 +183,14 @@ class SparseWeightExtras: } RETURN_TYPES = ("CN_WEIGHTS_EXTRAS", ) - RETURN_NAMES = ("extras", ) + RETURN_NAMES = ("cn_extras", ) FUNCTION = "create_weight_extras" CATEGORY = "Adv-ControlNet ๐Ÿ›‚๐Ÿ…๐Ÿ…’๐Ÿ…/SparseCtrl/extras" - def create_weight_extras(self, extras: dict[str]={}, sparse_hint_mult=1.0, sparse_nonhint_mult=1.0, sparse_mask_mult=1.0): - extras = extras.copy() - extras[SparseConst.HINT_MULT] = sparse_hint_mult - extras[SparseConst.NONHINT_MULT] = sparse_nonhint_mult - extras[SparseConst.MASK_MULT] = sparse_mask_mult - return (extras, ) + def create_weight_extras(self, cn_extras: dict[str]={}, sparse_hint_mult=1.0, sparse_nonhint_mult=1.0, sparse_mask_mult=1.0): + cn_extras = cn_extras.copy() + cn_extras[SparseConst.HINT_MULT] = sparse_hint_mult + cn_extras[SparseConst.NONHINT_MULT] = sparse_nonhint_mult + cn_extras[SparseConst.MASK_MULT] = sparse_mask_mult + return (cn_extras, ) diff --git a/adv_control/nodes_weight.py b/adv_control/nodes_weight.py index ddd7ef0..c7767db 100644 --- a/adv_control/nodes_weight.py +++ b/adv_control/nodes_weight.py @@ -12,7 +12,7 @@ class DefaultWeights: def INPUT_TYPES(s): return { "optional": { - "extras": ("CN_WEIGHTS_EXTRAS",), + "cn_extras": ("CN_WEIGHTS_EXTRAS",), } } @@ -22,8 +22,8 @@ class DefaultWeights: CATEGORY = "Adv-ControlNet ๐Ÿ›‚๐Ÿ…๐Ÿ…’๐Ÿ…/weights" - def load_weights(self, extras: dict[str]={}): - weights = ControlWeights.default(extras=extras) + def load_weights(self, cn_extras: dict[str]={}): + weights = ControlWeights.default(extras=cn_extras) return (weights, TimestepKeyframeGroup.default(TimestepKeyframe(control_weights=weights))) @@ -40,7 +40,7 @@ class ScaledSoftMaskedUniversalWeights: }, "optional": { "uncond_multiplier": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01}, ), - "extras": ("CN_WEIGHTS_EXTRAS",), + "cn_extras": ("CN_WEIGHTS_EXTRAS",), } } @@ -51,7 +51,7 @@ class ScaledSoftMaskedUniversalWeights: CATEGORY = "Adv-ControlNet ๐Ÿ›‚๐Ÿ…๐Ÿ…’๐Ÿ…/weights" def load_weights(self, mask: Tensor, min_base_multiplier: float, max_base_multiplier: float, lock_min=False, lock_max=False, - uncond_multiplier: float=1.0, extras: dict[str]={}): + uncond_multiplier: float=1.0, cn_extras: dict[str]={}): # normalize mask mask = mask.clone() x_min = 0.0 if lock_min else mask.min() @@ -60,7 +60,7 @@ class ScaledSoftMaskedUniversalWeights: mask = torch.ones_like(mask) * max_base_multiplier else: mask = linear_conversion(mask, x_min, x_max, min_base_multiplier, max_base_multiplier) - weights = ControlWeights.universal_mask(weight_mask=mask, uncond_multiplier=uncond_multiplier, extras=extras) + weights = ControlWeights.universal_mask(weight_mask=mask, uncond_multiplier=uncond_multiplier, extras=cn_extras) return (weights, TimestepKeyframeGroup.default(TimestepKeyframe(control_weights=weights))) @@ -74,7 +74,7 @@ class ScaledSoftUniversalWeights: }, "optional": { "uncond_multiplier": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01}, ), - "extras": ("CN_WEIGHTS_EXTRAS",), + "cn_extras": ("CN_WEIGHTS_EXTRAS",), } } @@ -84,8 +84,8 @@ class ScaledSoftUniversalWeights: CATEGORY = "Adv-ControlNet ๐Ÿ›‚๐Ÿ…๐Ÿ…’๐Ÿ…/weights" - def load_weights(self, base_multiplier, flip_weights, uncond_multiplier: float=1.0, extras: dict[str]={}): - weights = ControlWeights.universal(base_multiplier=base_multiplier, flip_weights=flip_weights, uncond_multiplier=uncond_multiplier, extras=extras) + def load_weights(self, base_multiplier, flip_weights, uncond_multiplier: float=1.0, cn_extras: dict[str]={}): + weights = ControlWeights.universal(base_multiplier=base_multiplier, flip_weights=flip_weights, uncond_multiplier=uncond_multiplier, extras=cn_extras) return (weights, TimestepKeyframeGroup.default(TimestepKeyframe(control_weights=weights))) @@ -111,7 +111,7 @@ class SoftControlNetWeights: }, "optional": { "uncond_multiplier": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01}, ), - "extras": ("CN_WEIGHTS_EXTRAS",), + "cn_extras": ("CN_WEIGHTS_EXTRAS",), } } @@ -123,10 +123,10 @@ class SoftControlNetWeights: def load_weights(self, weight_00, weight_01, weight_02, weight_03, weight_04, weight_05, weight_06, weight_07, weight_08, weight_09, weight_10, weight_11, weight_12, flip_weights, - uncond_multiplier: float=1.0, extras: dict[str]={}): + uncond_multiplier: float=1.0, cn_extras: dict[str]={}): weights = [weight_00, weight_01, weight_02, weight_03, weight_04, weight_05, weight_06, weight_07, weight_08, weight_09, weight_10, weight_11, weight_12] - weights = ControlWeights.controlnet(weights, flip_weights=flip_weights, uncond_multiplier=uncond_multiplier, extras=extras) + weights = ControlWeights.controlnet(weights, flip_weights=flip_weights, uncond_multiplier=uncond_multiplier, extras=cn_extras) return (weights, TimestepKeyframeGroup.default(TimestepKeyframe(control_weights=weights))) @@ -152,7 +152,7 @@ class CustomControlNetWeights: }, "optional": { "uncond_multiplier": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01}, ), - "extras": ("CN_WEIGHTS_EXTRAS",), + "cn_extras": ("CN_WEIGHTS_EXTRAS",), } } @@ -164,10 +164,10 @@ class CustomControlNetWeights: def load_weights(self, weight_00, weight_01, weight_02, weight_03, weight_04, weight_05, weight_06, weight_07, weight_08, weight_09, weight_10, weight_11, weight_12, flip_weights, - uncond_multiplier: float=1.0, extras: dict[str]={}): + uncond_multiplier: float=1.0, cn_extras: dict[str]={}): weights = [weight_00, weight_01, weight_02, weight_03, weight_04, weight_05, weight_06, weight_07, weight_08, weight_09, weight_10, weight_11, weight_12] - weights = ControlWeights.controlnet(weights, flip_weights=flip_weights, uncond_multiplier=uncond_multiplier, extras=extras) + weights = ControlWeights.controlnet(weights, flip_weights=flip_weights, uncond_multiplier=uncond_multiplier, extras=cn_extras) return (weights, TimestepKeyframeGroup.default(TimestepKeyframe(control_weights=weights))) @@ -184,7 +184,7 @@ class SoftT2IAdapterWeights: }, "optional": { "uncond_multiplier": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01}, ), - "extras": ("CN_WEIGHTS_EXTRAS",), + "cn_extras": ("CN_WEIGHTS_EXTRAS",), } } @@ -195,10 +195,10 @@ class SoftT2IAdapterWeights: CATEGORY = "Adv-ControlNet ๐Ÿ›‚๐Ÿ…๐Ÿ…’๐Ÿ…/weights/T2IAdapter" def load_weights(self, weight_00, weight_01, weight_02, weight_03, flip_weights, - uncond_multiplier: float=1.0, extras: dict[str]={}): + uncond_multiplier: float=1.0, cn_extras: dict[str]={}): weights = [weight_00, weight_01, weight_02, weight_03] weights = get_properly_arranged_t2i_weights(weights) - weights = ControlWeights.t2iadapter(weights, flip_weights=flip_weights, uncond_multiplier=uncond_multiplier, extras=extras) + weights = ControlWeights.t2iadapter(weights, flip_weights=flip_weights, uncond_multiplier=uncond_multiplier, extras=cn_extras) return (weights, TimestepKeyframeGroup.default(TimestepKeyframe(control_weights=weights))) @@ -215,7 +215,7 @@ class CustomT2IAdapterWeights: }, "optional": { "uncond_multiplier": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01}, ), - "extras": ("CN_WEIGHTS_EXTRAS",), + "cn_extras": ("CN_WEIGHTS_EXTRAS",), } } @@ -226,8 +226,8 @@ class CustomT2IAdapterWeights: CATEGORY = "Adv-ControlNet ๐Ÿ›‚๐Ÿ…๐Ÿ…’๐Ÿ…/weights/T2IAdapter" def load_weights(self, weight_00, weight_01, weight_02, weight_03, flip_weights, - uncond_multiplier: float=1.0, extras: dict[str]={}): + uncond_multiplier: float=1.0, cn_extras: dict[str]={}): weights = [weight_00, weight_01, weight_02, weight_03] weights = get_properly_arranged_t2i_weights(weights) - weights = ControlWeights.t2iadapter(weights, flip_weights=flip_weights, uncond_multiplier=uncond_multiplier, extras=extras) + weights = ControlWeights.t2iadapter(weights, flip_weights=flip_weights, uncond_multiplier=uncond_multiplier, extras=cn_extras) return (weights, TimestepKeyframeGroup.default(TimestepKeyframe(control_weights=weights)))