diff --git a/control.py b/control.py index d3d4c78..d9dc96f 100644 --- a/control.py +++ b/control.py @@ -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: diff --git a/nodes.py b/nodes.py index 849f7b1..d5ceee8 100644 --- a/nodes.py +++ b/nodes.py @@ -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" }