Updated ControlNetAdvanced to work with Flux controlnet, beginning of rework of weights

This commit is contained in:
Jedrzej Kosinski
2024-08-24 00:43:21 -05:00
parent 949843e2c0
commit 1807cd64e4
3 changed files with 24 additions and 15 deletions
+15 -9
View File
@@ -8,7 +8,7 @@ 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
from comfy.controlnet import ControlBase, ControlNet, ControlLora, T2IAdapter, StrengthType
from comfy.model_patcher import ModelPatcher
from .control_sparsectrl import SparseModelPatcher, SparseControlNet, SparseCtrlMotionWrapper, SparseSettings, SparseConst
@@ -21,7 +21,7 @@ from .logger import logger
class ControlNetAdvanced(ControlNet, AdvancedControlBase):
def __init__(self, control_model, timestep_keyframes: TimestepKeyframeGroup, global_average_pooling=False, compression_ratio=8, latent_format=None, device=None, load_device=None, manual_cast_dtype=None):
def __init__(self, control_model, timestep_keyframes: TimestepKeyframeGroup, global_average_pooling=False, compression_ratio=8, latent_format=None, device=None, load_device=None, manual_cast_dtype=None, extra_conds=["y"], strength_type=StrengthType.CONSTANT):
super().__init__(control_model=control_model, global_average_pooling=global_average_pooling, compression_ratio=compression_ratio, latent_format=latent_format, device=device, load_device=load_device, manual_cast_dtype=manual_cast_dtype)
AdvancedControlBase.__init__(self, super(), timestep_keyframes=timestep_keyframes, weights_default=ControlWeights.controlnet())
@@ -64,9 +64,9 @@ class ControlNetAdvanced(ControlNet, AdvancedControlBase):
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] * compression_ratio, x_noisy.shape[2] * compression_ratio, 'nearest-exact', "center")
self.cond_hint = comfy.utils.common_upscale(actual_cond_hint_orig[self.sub_idxs], x_noisy.shape[3] * compression_ratio, x_noisy.shape[2] * compression_ratio, self.upscale_algorithm, "center")
else:
self.cond_hint = comfy.utils.common_upscale(self.cond_hint_original, x_noisy.shape[3] * compression_ratio, x_noisy.shape[2] * compression_ratio, 'nearest-exact', "center")
self.cond_hint = comfy.utils.common_upscale(self.cond_hint_original, x_noisy.shape[3] * compression_ratio, x_noisy.shape[2] * compression_ratio, self.upscale_algorithm, "center")
if self.vae is not None:
loaded_models = comfy.model_management.loaded_models(only_currently_used=True)
self.cond_hint = self.vae.encode(self.cond_hint.movedim(1, -1))
@@ -81,17 +81,22 @@ class ControlNetAdvanced(ControlNet, AdvancedControlBase):
self.prepare_mask_cond_hint(x_noisy=x_noisy, t=t, cond=cond, batched_number=batched_number, dtype=dtype)
context = cond.get('crossattn_controlnet', cond['c_crossattn'])
y = cond.get('y', None)
if y is not None:
y = y.to(dtype)
extra = self.extra_args.copy()
for c in self.extra_conds:
temp = cond.get(c, None)
if temp is not None:
extra[c] = temp.to(dtype)
timestep = self.model_sampling_current.timestep(t)
x_noisy = self.model_sampling_current.calculate_input(t, x_noisy)
control = self.control_model(x=x_noisy.to(dtype), hint=self.cond_hint, timesteps=timestep.float(), context=context.to(dtype), y=y)
control = self.control_model(x=x_noisy.to(dtype), hint=self.cond_hint, timesteps=timestep.to(dtype), context=context.to(dtype), **extra)
return self.control_merge(control, control_prev, output_dtype)
def copy(self):
c = ControlNetAdvanced(self.control_model, self.timestep_keyframes, global_average_pooling=self.global_average_pooling, load_device=self.load_device, manual_cast_dtype=self.manual_cast_dtype)
c.control_model = self.control_model
c.control_model_wrapped = self.control_model_wrapped
self.copy_to(c)
self.copy_to_advanced(c)
return c
@@ -99,7 +104,8 @@ class ControlNetAdvanced(ControlNet, AdvancedControlBase):
@staticmethod
def from_vanilla(v: ControlNet, timestep_keyframe: TimestepKeyframeGroup=None) -> 'ControlNetAdvanced':
to_return = ControlNetAdvanced(control_model=v.control_model, timestep_keyframes=timestep_keyframe,
global_average_pooling=v.global_average_pooling, compression_ratio=v.compression_ratio, latent_format=v.latent_format, device=v.device, load_device=v.load_device, manual_cast_dtype=v.manual_cast_dtype)
global_average_pooling=v.global_average_pooling, compression_ratio=v.compression_ratio, latent_format=v.latent_format, device=v.device, load_device=v.load_device,
manual_cast_dtype=v.manual_cast_dtype)
v.copy_to(to_return)
return to_return
+6 -6
View File
@@ -42,7 +42,7 @@ class ScaledSoftMaskedUniversalWeights:
"optional": {
"uncond_multiplier": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01}, ),
"cn_extras": ("CN_WEIGHTS_EXTRAS",),
"autosize": ("ACNAUTOSIZE", {"padding": 75}),
"autosize": ("ACNAUTOSIZE", {"padding": 0}),
}
}
@@ -77,7 +77,7 @@ class ScaledSoftUniversalWeights:
"optional": {
"uncond_multiplier": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01}, ),
"cn_extras": ("CN_WEIGHTS_EXTRAS",),
"autosize": ("ACNAUTOSIZE", {"padding": 20}),
"autosize": ("ACNAUTOSIZE", {"padding": 0}),
}
}
@@ -115,7 +115,7 @@ class SoftControlNetWeights:
"optional": {
"uncond_multiplier": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01}, ),
"cn_extras": ("CN_WEIGHTS_EXTRAS",),
"autosize": ("ACNAUTOSIZE", {"padding": 50}),
"autosize": ("ACNAUTOSIZE", {"padding": 0}),
}
}
@@ -157,7 +157,7 @@ class CustomControlNetWeights:
"optional": {
"uncond_multiplier": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01}, ),
"cn_extras": ("CN_WEIGHTS_EXTRAS",),
"autosize": ("ACNAUTOSIZE", {"padding": 65}),
"autosize": ("ACNAUTOSIZE", {"padding": 0}),
}
}
@@ -190,7 +190,7 @@ class SoftT2IAdapterWeights:
"optional": {
"uncond_multiplier": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01}, ),
"cn_extras": ("CN_WEIGHTS_EXTRAS",),
"autosize": ("ACNAUTOSIZE", {"padding": 55}),
"autosize": ("ACNAUTOSIZE", {"padding": 0}),
}
}
@@ -223,7 +223,7 @@ class CustomT2IAdapterWeights:
"optional": {
"uncond_multiplier": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01}, ),
"cn_extras": ("CN_WEIGHTS_EXTRAS",),
"autosize": ("ACNAUTOSIZE", {"padding": 65}),
"autosize": ("ACNAUTOSIZE", {"padding": 0}),
}
}
+3
View File
@@ -175,6 +175,9 @@ class ControlWeights:
def get(self, idx: int, control: dict[str, list[Tensor]], key: str, default=1.0) -> Union[float, Tensor]:
# if weights is not none, return index
if self.weights is not None:
# support live calculation
if callable(self.weights):
return self.weights(idx=idx, control=control, key=key)
# if middle weight, need to pretend index is actually after all the output weights (if applicable)
if key == "middle" and "output" in control:
idx += len(control["output"])