nodes/Kosinkadink-ComfyUI-Advanced-ControlNet
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Merge PR #116 from Kosinkadink/develop - SparseCtrl upgrades + bugfixes

Browse Source 
SparseCtrl Upgrades + bug fixes
This commit is contained in:
Jedrzej Kosinski
2024-05-29 01:43:03 -05:00
committed by GitHub
parent 67884487aa 576426a734
commit 95441b6d7d
8 changed files with 328 additions and 126 deletions
Download Patch File Download Diff File Expand all files Collapse all files
adv_control/control.py
+65 -26
View File
@@ -3,18 +3,20 @@ from torch import Tensor
import torch
import os
import comfy.ops
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_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,
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
@@ -55,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)
@@ -97,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]
@@ -131,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)
@@ -221,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)
@@ -291,18 +302,33 @@ 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.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 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))
# 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)))
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)
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
@@ -319,7 +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] = 1.0
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)
@@ -327,7 +353,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)
@@ -342,6 +368,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:
@@ -359,6 +391,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)
@@ -411,12 +445,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
@@ -551,13 +588,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 +714,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)
adv_control/control_reference.py
+2 -2
View File
@@ -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)
adv_control/control_sparsectrl.py
+161 -54
View File
@@ -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
@@ -24,22 +25,28 @@ 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 .logger import logger
from .utils import (BIGMAX, 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
# 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
@@ -47,6 +54,12 @@ else:
optimized_attention_mm = attention_sub_quad
class SparseConst:
HINT_MULT = "sparse_hint_mult"
NONHINT_MULT = "sparse_nonhint_mult"
MASK_MULT = "sparse_mask_mult"
class SparseControlNet(ControlNetCLDM):
def __init__(self, *args,**kwargs):
super().__init__(*args, **kwargs)
@@ -100,28 +113,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
@@ -173,14 +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):
def __init__(self, sparse_method: 'SparseMethod', use_motion: bool=True, motion_strength=1.0, motion_scale=1.0, merged=False,
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_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)
@@ -193,9 +220,62 @@ 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) -> list[int]:
pass
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])
class SparseSpreadMethod(SparseMethod):
UNIFORM = "uniform"
@@ -209,7 +289,7 @@ 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) -> list[int]:
# if hint_length >= full_length, limit hints to full_length
if hint_length >= full_length:
return list(range(full_length))
@@ -247,7 +327,7 @@ class SparseIndexMethod(SparseMethod):
super().__init__(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
@@ -318,7 +398,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
@@ -328,13 +408,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
@@ -448,22 +528,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:
@@ -490,9 +570,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):
@@ -507,6 +587,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
@@ -521,6 +602,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:
@@ -552,8 +634,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)
@@ -578,6 +660,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
@@ -592,10 +675,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(
[
@@ -613,11 +696,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
@@ -641,6 +725,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
@@ -661,10 +747,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,
@@ -739,6 +825,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__()
@@ -761,15 +848,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:
@@ -779,6 +867,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 +925,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):
operations=disable_weight_init_clean_groupnorm):
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 +943,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,21 +964,29 @@ 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,
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
@@ -911,6 +1015,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,
adv_control/nodes.py
+4 -2
View File
@@ -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 🛂🅐🅒🅝",
adv_control/nodes_sparsectrl.py
+36 -3
View File
@@ -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, SparseContextAware
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": (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_mult": ("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=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)
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_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,)
@@ -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": {
"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}, ),
}
}
RETURN_TYPES = ("CN_WEIGHTS_EXTRAS", )
RETURN_NAMES = ("cn_extras", )
FUNCTION = "create_weight_extras"
CATEGORY = "Adv-ControlNet 🛂🅐🅒🅝/SparseCtrl/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, )
adv_control/nodes_weight.py
+23 -14
View File
@@ -11,6 +11,9 @@ class DefaultWeights:
@classmethod
def INPUT_TYPES(s):
return {
"optional": {
"cn_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, cn_extras: dict[str]={}):
weights = ControlWeights.default(extras=cn_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}, ),
"cn_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, cn_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=cn_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}, ),
"cn_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, 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)))
@@ -106,6 +111,7 @@ class SoftControlNetWeights:
},
"optional": {
"uncond_multiplier": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01}, ),
"cn_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, 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)
weights = ControlWeights.controlnet(weights, flip_weights=flip_weights, uncond_multiplier=uncond_multiplier, extras=cn_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}, ),
"cn_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, 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)
weights = ControlWeights.controlnet(weights, flip_weights=flip_weights, uncond_multiplier=uncond_multiplier, extras=cn_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}, ),
"cn_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, 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)
weights = ControlWeights.t2iadapter(weights, flip_weights=flip_weights, uncond_multiplier=uncond_multiplier, extras=cn_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}, ),
"cn_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, 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)
weights = ControlWeights.t2iadapter(weights, flip_weights=flip_weights, uncond_multiplier=uncond_multiplier, extras=cn_extras)
return (weights, TimestepKeyframeGroup.default(TimestepKeyframe(control_weights=weights)))
adv_control/utils.py
+36 -24
View File
@@ -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, extras: dict[str]={},):
self.weight_type = weight_type
self.base_multiplier = base_multiplier
self.flip_weights = flip_weights
@@ -162,6 +162,9 @@ 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
self.extras = extras
def get(self, idx: int, default=1.0) -> Union[float, Tensor]:
# if weights is not none, return index
@@ -174,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:
@@ -433,8 +436,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 +454,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:
@@ -523,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
@@ -576,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)
@@ -643,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
@@ -772,6 +782,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 +872,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
pyproject.toml
+1 -1
View File
@@ -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 = []