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

updated to work with latest ComfyUI, added Load Images node to batch load images (sorted), added timestep_keyframe output on weight nodes to simplify node connections if not chaining timestep_keyframes together

Browse Source
This commit is contained in:
Jedrzej Kosinski
2023-08-26 12:43:05 -05:00
parent 354a8ce10c
commit 25b60aa491
2 changed files with 88 additions and 23 deletions
Download Patch File Download Diff File Expand all files Collapse all files
control.py
+19 -12
View File
@@ -14,7 +14,8 @@ sys.path.insert(0, os.path.join(os.path.dirname(os.path.realpath(__file__)), "co
from comfy.cldm import cldm
from comfy.t2i_adapter import adapter
from comfy.sd import ControlBase, ModelPatcher, broadcast_image_to, ControlLora
from comfy.sd import ModelPatcher
from comfy.controlnet import ControlBase, broadcast_image_to, ControlLora
import comfy.utils as utils
import comfy.model_management as model_management
import comfy.model_detection as model_detection
@@ -61,7 +62,7 @@ class LatentKeyframeGroup:
class TimestepKeyframe:
def __init__(self,
start_percent: float,
start_percent: float = 0.0,
control_net_weights: ControlNetWeightsType = None,
t2i_adapter_weights: T2IAdapterWeightsType = None,
latent_keyframes: LatentKeyframeGroup = None) -> None:
@@ -105,6 +106,12 @@ class TimestepKeyframeGroup:
def is_empty(self) -> bool:
return len(self.keyframes) == 0
@classmethod
def default(cls, keyframe: TimestepKeyframe) -> 'TimestepKeyframeGroup':
group = cls()
group.keyframes[0] = keyframe
return group
# Copied from comfy.sd, weights modified
@@ -113,7 +120,6 @@ class ControlNetAdvanced(ControlBase):
super().__init__(device)
self.control_model = control_model
self.control_model_wrapped = ModelPatcher(self.control_model, load_device=model_management.get_torch_device(), offload_device=model_management.unet_offload_device())
self.timestep_keyframes = timestep_keyframes if timestep_keyframes else TimestepKeyframeGroup()
self.weights = self.timestep_keyframes.keyframes[0].control_net_weights if self.timestep_keyframes.keyframes[0].control_net_weights else [1.0]*13
@@ -166,16 +172,17 @@ class ControlNetAdvanced(ControlBase):
if self.global_average_pooling:
x = torch.mean(x, dim=(2, 3), keepdim=True).repeat(1, 1, x.shape[2], x.shape[3])
# get batch indeces to zero out, AKA latents that should not be influenced by ControlNet
indeces_to_zero = set(range(x.size()[0]//2))
for keyframe in current_timestep_keyframe.latent_keyframes:
if keyframe.batch_index in indeces_to_zero:
indeces_to_zero.remove(keyframe.batch_index)
if current_timestep_keyframe.latent_keyframes is not None:
# get batch indeces to zero out, AKA latents that should not be influenced by ControlNet
indeces_to_zero = set(range(x.size()[0]//2))
for keyframe in current_timestep_keyframe.latent_keyframes:
if keyframe.batch_index in indeces_to_zero:
indeces_to_zero.remove(keyframe.batch_index)
# zero them out by multiplying by zero
for batch_index in indeces_to_zero:
x[batch_index] *= 0.0
x[(x.size()[0]//2) + batch_index] *= 0.0
# zero them out by multiplying by zero
for batch_index in indeces_to_zero:
x[batch_index] *= 0.0
x[(x.size()[0]//2) + batch_index] *= 0.0
x *= self.strength * self.weights[i]
if x.dtype != output_dtype and not autocast_enabled:
nodes.py
+69 -11
View File
@@ -1,10 +1,14 @@
import sys
import os
import torch
import numpy as np
from PIL import Image, ImageOps
import folder_paths
sys.path.insert(0, os.path.join(os.path.dirname(os.path.realpath(__file__)), "comfy"))
from comfy.sd import ControlBase
from .control import load_controlnet, ControlNetWeightsType, T2IAdapterWeightsType,\
LatentKeyframe, LatentKeyframeGroup, TimestepKeyframe, TimestepKeyframeGroup
@@ -20,7 +24,7 @@ class ScaledSoftControlNetWeights:
},
}
RETURN_TYPES = ("CONTROL_NET_WEIGHTS", )
RETURN_TYPES = ("CONTROL_NET_WEIGHTS", "TIMESTEP_KEYFRAME",)
FUNCTION = "load_weights"
CATEGORY = "adv-controlnet/weights"
@@ -29,7 +33,7 @@ class ScaledSoftControlNetWeights:
weights = [(base_multiplier ** float(12 - i)) for i in range(13)]
if flip_weights:
weights.reverse()
return (weights, )
return (weights, TimestepKeyframeGroup.default(TimestepKeyframe(control_net_weights=weights)))
class SoftControlNetWeights:
@@ -54,7 +58,7 @@ class SoftControlNetWeights:
},
}
RETURN_TYPES = ("CONTROL_NET_WEIGHTS", )
RETURN_TYPES = ("CONTROL_NET_WEIGHTS", "TIMESTEP_KEYFRAME",)
FUNCTION = "load_weights"
CATEGORY = "adv-controlnet/weights"
@@ -65,7 +69,7 @@ class SoftControlNetWeights:
weight_07, weight_08, weight_09, weight_10, weight_11, weight_12]
if flip_weights:
weights.reverse()
return (weights,)
return (weights, TimestepKeyframeGroup.default(TimestepKeyframe(control_net_weights=weights)))
class CustomControlNetWeights:
@@ -90,7 +94,7 @@ class CustomControlNetWeights:
}
}
RETURN_TYPES = ("CONTROL_NET_WEIGHTS", )
RETURN_TYPES = ("CONTROL_NET_WEIGHTS", "TIMESTEP_KEYFRAME",)
FUNCTION = "load_weights"
CATEGORY = "adv-controlnet/weights"
@@ -101,7 +105,7 @@ class CustomControlNetWeights:
weight_07, weight_08, weight_09, weight_10, weight_11, weight_12]
if flip_weights:
weights.reverse()
return (weights,)
return (weights, TimestepKeyframeGroup.default(TimestepKeyframe(control_net_weights=weights)))
class SoftT2IAdapterWeights:
@@ -117,7 +121,7 @@ class SoftT2IAdapterWeights:
},
}
RETURN_TYPES = ("T2I_ADAPTER_WEIGHTS", )
RETURN_TYPES = ("T2I_ADAPTER_WEIGHTS", "TIMESTEP_KEYFRAME",)
FUNCTION = "load_weights"
CATEGORY = "adv-controlnet/weights"
@@ -126,7 +130,7 @@ class SoftT2IAdapterWeights:
weights = [weight_00, weight_01, weight_02, weight_03]
if flip_weights:
weights.reverse()
return (weights,)
return (weights, TimestepKeyframeGroup.default(TimestepKeyframe(t2i_adapter_weights=weights)))
class CustomT2IAdapterWeights:
@@ -142,7 +146,7 @@ class CustomT2IAdapterWeights:
},
}
RETURN_TYPES = ("T2I_ADAPTER_WEIGHTS", )
RETURN_TYPES = ("T2I_ADAPTER_WEIGHTS", "TIMESTEP_KEYFRAME",)
FUNCTION = "load_weights"
CATEGORY = "adv-controlnet/weights"
@@ -151,7 +155,8 @@ class CustomT2IAdapterWeights:
weights = [weight_00, weight_01, weight_02, weight_03]
if flip_weights:
weights.reverse()
return (weights,)
return (weights, TimestepKeyframeGroup.default(TimestepKeyframe(t2i_adapter_weights=weights)))
class TimestepKeyframeNode:
@@ -373,6 +378,55 @@ class ControlNetApplyPartialBatch: # NOT USED: was used for a different test, ha
return (out[0], out[1])
class LoadImagesFromDirectory:
@classmethod
def INPUT_TYPES(s):
return {
"required": {
"directory": ("STRING", {"default": ""}),
}
}
RETURN_TYPES = ("IMAGE", "MASK")
FUNCTION = "load_images"
CATEGORY = "adv-controlnet/image"
def load_images(self, directory):
if not os.path.isdir(directory):
raise FileNotFoundError(f"Directory '{directory} cannot be found.'")
dir_files = os.listdir(directory)
if len(dir_files) == 0:
raise FileNotFoundError(f"No files in directory '{directory}'.")
dir_files = sorted(dir_files)
dir_files = [os.path.join(directory, x) for x in dir_files]
images = []
masks = []
for image_path in dir_files:
i = Image.open(image_path)
i = ImageOps.exif_transpose(i)
image = i.convert("RGB")
image = np.array(image).astype(np.float32) / 255.0
image = torch.from_numpy(image)[None,]
if 'A' in i.getbands():
mask = np.array(i.getchannel('A')).astype(np.float32) / 255.0
mask = 1. - torch.from_numpy(mask)
else:
mask = torch.zeros((64,64), dtype=torch.float32, device="cpu")
images.append(image)
masks.append(mask)
if len(images) == 0:
raise FileNotFoundError(f"No images could be loaded from directory '{directory}'.")
return (torch.cat(images, dim=0), torch.cat(masks, dim=0))
# NODE MAPPING
NODE_CLASS_MAPPINGS = {
# Keyframes
@@ -389,6 +443,8 @@ NODE_CLASS_MAPPINGS = {
"CustomControlNetWeights": CustomControlNetWeights,
"SoftT2IAdapterWeights": SoftT2IAdapterWeights,
"CustomT2IAdapterWeights": CustomT2IAdapterWeights,
# Image
"LoadImagesFromDirectory": LoadImagesFromDirectory
}
NODE_DISPLAY_NAME_MAPPINGS = {
@@ -406,4 +462,6 @@ NODE_DISPLAY_NAME_MAPPINGS = {
"CustomControlNetWeights": "Custom ControlNet Weights",
"SoftT2IAdapterWeights": "Soft T2IAdapter Weights",
"CustomT2IAdapterWeights": "Custom T2IAdapter Weights",
# Image
"LoadImagesFromDirectory": "Load Images"
}