From 3e413e673609279ce4704e99932e17e45f57ea37 Mon Sep 17 00:00:00 2001 From: Alschain Date: Fri, 30 Aug 2024 21:13:03 +0800 Subject: [PATCH] feat: custom function support --- ComfyUI-Manager | 1 + ComfyUI-to-Python-Extension | 1 + comfy/ldm/flux/controlnet_xlabs.py | 104 ++++++++++++++++ custom_functions/image_text_matting.py | 151 +++++++++++++++++++++++ custom_functions/image_text_to_image.py | 156 ++++++++++++++++++++++++ nodes.py | 7 +- 6 files changed, 418 insertions(+), 2 deletions(-) create mode 160000 ComfyUI-Manager create mode 160000 ComfyUI-to-Python-Extension create mode 100644 comfy/ldm/flux/controlnet_xlabs.py create mode 100644 custom_functions/image_text_matting.py create mode 100644 custom_functions/image_text_to_image.py diff --git a/ComfyUI-Manager b/ComfyUI-Manager new file mode 160000 index 000000000..d7170c026 --- /dev/null +++ b/ComfyUI-Manager @@ -0,0 +1 @@ +Subproject commit d7170c0264ca1623b7a4ba7d63756a9b54a7d119 diff --git a/ComfyUI-to-Python-Extension b/ComfyUI-to-Python-Extension new file mode 160000 index 000000000..d3b82726a --- /dev/null +++ b/ComfyUI-to-Python-Extension @@ -0,0 +1 @@ +Subproject commit d3b82726a7f185f25f3d0194702ea17df47ee167 diff --git a/comfy/ldm/flux/controlnet_xlabs.py b/comfy/ldm/flux/controlnet_xlabs.py new file mode 100644 index 000000000..5d700f16c --- /dev/null +++ b/comfy/ldm/flux/controlnet_xlabs.py @@ -0,0 +1,104 @@ +#Original code can be found on: https://github.com/XLabs-AI/x-flux/blob/main/src/flux/controlnet.py + +import torch +from torch import Tensor, nn +from einops import rearrange, repeat + +from .layers import (DoubleStreamBlock, EmbedND, LastLayer, + MLPEmbedder, SingleStreamBlock, + timestep_embedding) + +from .model import Flux +import comfy.ldm.common_dit + + +class ControlNetFlux(Flux): + def __init__(self, image_model=None, dtype=None, device=None, operations=None, **kwargs): + super().__init__(final_layer=False, dtype=dtype, device=device, operations=operations, **kwargs) + + # add ControlNet blocks + self.controlnet_blocks = nn.ModuleList([]) + for _ in range(self.params.depth): + controlnet_block = operations.Linear(self.hidden_size, self.hidden_size, dtype=dtype, device=device) + # controlnet_block = zero_module(controlnet_block) + self.controlnet_blocks.append(controlnet_block) + self.pos_embed_input = operations.Linear(self.in_channels, self.hidden_size, bias=True, dtype=dtype, device=device) + self.gradient_checkpointing = False + self.input_hint_block = nn.Sequential( + operations.Conv2d(3, 16, 3, padding=1, dtype=dtype, device=device), + nn.SiLU(), + operations.Conv2d(16, 16, 3, padding=1, dtype=dtype, device=device), + nn.SiLU(), + operations.Conv2d(16, 16, 3, padding=1, stride=2, dtype=dtype, device=device), + nn.SiLU(), + operations.Conv2d(16, 16, 3, padding=1, dtype=dtype, device=device), + nn.SiLU(), + operations.Conv2d(16, 16, 3, padding=1, stride=2, dtype=dtype, device=device), + nn.SiLU(), + operations.Conv2d(16, 16, 3, padding=1, dtype=dtype, device=device), + nn.SiLU(), + operations.Conv2d(16, 16, 3, padding=1, stride=2, dtype=dtype, device=device), + nn.SiLU(), + operations.Conv2d(16, 16, 3, padding=1, dtype=dtype, device=device) + ) + + def forward_orig( + self, + img: Tensor, + img_ids: Tensor, + controlnet_cond: Tensor, + txt: Tensor, + txt_ids: Tensor, + timesteps: Tensor, + y: Tensor, + guidance: Tensor = None, + ) -> Tensor: + if img.ndim != 3 or txt.ndim != 3: + raise ValueError("Input img and txt tensors must have 3 dimensions.") + + # running on sequences img + img = self.img_in(img) + controlnet_cond = self.input_hint_block(controlnet_cond) + controlnet_cond = rearrange(controlnet_cond, "b c (h ph) (w pw) -> b (h w) (c ph pw)", ph=2, pw=2) + controlnet_cond = self.pos_embed_input(controlnet_cond) + img = img + controlnet_cond + vec = self.time_in(timestep_embedding(timesteps, 256)) + if self.params.guidance_embed: + vec = vec + self.guidance_in(timestep_embedding(guidance, 256)) + vec = vec + self.vector_in(y) + txt = self.txt_in(txt) + + ids = torch.cat((txt_ids, img_ids), dim=1) + pe = self.pe_embedder(ids) + + block_res_samples = () + + for block in self.double_blocks: + img, txt = block(img=img, txt=txt, vec=vec, pe=pe) + block_res_samples = block_res_samples + (img,) + + controlnet_block_res_samples = () + for block_res_sample, controlnet_block in zip(block_res_samples, self.controlnet_blocks): + block_res_sample = controlnet_block(block_res_sample) + controlnet_block_res_samples = controlnet_block_res_samples + (block_res_sample,) + + return {"input": (controlnet_block_res_samples * 10)[:19]} + + def forward(self, x, timesteps, context, y, guidance=None, hint=None, **kwargs): + hint = hint * 2.0 - 1.0 + + bs, c, h, w = x.shape + patch_size = 2 + x = comfy.ldm.common_dit.pad_to_patch_size(x, (patch_size, patch_size)) + + img = rearrange(x, "b c (h ph) (w pw) -> b (h w) (c ph pw)", ph=patch_size, pw=patch_size) + + h_len = ((h + (patch_size // 2)) // patch_size) + w_len = ((w + (patch_size // 2)) // patch_size) + img_ids = torch.zeros((h_len, w_len, 3), device=x.device, dtype=x.dtype) + img_ids[..., 1] = img_ids[..., 1] + torch.linspace(0, h_len - 1, steps=h_len, device=x.device, dtype=x.dtype)[:, None] + img_ids[..., 2] = img_ids[..., 2] + torch.linspace(0, w_len - 1, steps=w_len, device=x.device, dtype=x.dtype)[None, :] + img_ids = repeat(img_ids, "h w c -> b (h w) c", b=bs) + + txt_ids = torch.zeros((bs, context.shape[1], 3), device=x.device, dtype=x.dtype) + return self.forward_orig(img, img_ids, hint, context, txt_ids, timesteps, y, guidance) diff --git a/custom_functions/image_text_matting.py b/custom_functions/image_text_matting.py new file mode 100644 index 000000000..d87fdcc8e --- /dev/null +++ b/custom_functions/image_text_matting.py @@ -0,0 +1,151 @@ +import os +import sys +from typing import Sequence, Mapping, Any, Union +import torch + + +def get_value_at_index(obj: Union[Sequence, Mapping], index: int) -> Any: + """Returns the value at the given index of a sequence or mapping. + + If the object is a sequence (like list or string), returns the value at the given index. + If the object is a mapping (like a dictionary), returns the value at the index-th key. + + Some return a dictionary, in these cases, we look for the "results" key + + Args: + obj (Union[Sequence, Mapping]): The object to retrieve the value from. + index (int): The index of the value to retrieve. + + Returns: + Any: The value at the given index. + + Raises: + IndexError: If the index is out of bounds for the object and the object is not a mapping. + """ + try: + return obj[index] + except KeyError: + return obj["result"][index] + + +def find_path(name: str, path: str = None) -> str: + """ + Recursively looks at parent folders starting from the given path until it finds the given name. + Returns the path as a Path object if found, or None otherwise. + """ + # If no path is given, use the current working directory + if path is None: + path = os.getcwd() + + # Check if the current directory contains the name + if name in os.listdir(path): + path_name = os.path.join(path, name) + print(f"{name} found: {path_name}") + return path_name + + # Get the parent directory + parent_directory = os.path.dirname(path) + + # If the parent directory is the same as the current directory, we've reached the root and stop the search + if parent_directory == path: + return None + + # Recursively call the function with the parent directory + return find_path(name, parent_directory) + + +def add_comfyui_directory_to_sys_path() -> None: + """ + Add 'ComfyUI' to the sys.path + """ + comfyui_path = find_path("ComfyUI") + if comfyui_path is not None and os.path.isdir(comfyui_path): + sys.path.append(comfyui_path) + print(f"'{comfyui_path}' added to sys.path") + + +def add_extra_model_paths() -> None: + """ + Parse the optional extra_model_paths.yaml file and add the parsed paths to the sys.path. + """ + from main import load_extra_path_config + + extra_model_paths = find_path("extra_model_paths.yaml") + + if extra_model_paths is not None: + load_extra_path_config(extra_model_paths) + else: + print("Could not find the extra_model_paths config file.") + + +from nodes import NODE_CLASS_MAPPINGS, LoadImage, init_extra_nodes + + +def image_text_matting(image_path, text, abs_path=True): + add_comfyui_directory_to_sys_path() + add_extra_model_paths() + init_extra_nodes(True) + + with torch.inference_mode(): + sammodelloader_segment_anything = NODE_CLASS_MAPPINGS[ + "SAMModelLoader (segment anything)" + ]() + sammodelloader_segment_anything_1 = sammodelloader_segment_anything.main( + model_name="sam_hq_vit_h (2.57GB)" + ) + + groundingdinomodelloader_segment_anything = NODE_CLASS_MAPPINGS[ + "GroundingDinoModelLoader (segment anything)" + ]() + groundingdinomodelloader_segment_anything_2 = ( + groundingdinomodelloader_segment_anything.main( + model_name="GroundingDINO_SwinT_OGC (694MB)" + ) + ) + + mattingmodelloader = NODE_CLASS_MAPPINGS["MattingModelLoader"]() + mattingmodelloader_8 = mattingmodelloader.main( + model_name="vitmatte_small (103 MB)" + ) + + loadimage = LoadImage() + loadimage_12 = loadimage.load_image(image=image_path, abs_path=abs_path) + + groundingdinosamsegment_segment_anything = NODE_CLASS_MAPPINGS[ + "GroundingDinoSAMSegment (segment anything)" + ]() + createtrimap = NODE_CLASS_MAPPINGS["CreateTrimap"]() + applymatting = NODE_CLASS_MAPPINGS["ApplyMatting"]() + + groundingdinosamsegment_segment_anything_3 = ( + groundingdinosamsegment_segment_anything.main( + prompt=text, + threshold=0.3, + sam_model=get_value_at_index(sammodelloader_segment_anything_1, 0), + grounding_dino_model=get_value_at_index( + groundingdinomodelloader_segment_anything_2, 0 + ), + image=get_value_at_index(loadimage_12, 0), + ) + ) + + createtrimap_11 = createtrimap.main( + kernel_size=20.86, + mask=get_value_at_index(groundingdinosamsegment_segment_anything_3, 1), + ) + + + applymatting_9 = applymatting.main( + matting_model=get_value_at_index(mattingmodelloader_8, 0), + matting_preprocessor=get_value_at_index(mattingmodelloader_8, 1), + image=get_value_at_index(loadimage_12, 0), + trimap=get_value_at_index(createtrimap_11, 0), + ) + + output_image = get_value_at_index(applymatting_9, 1) + return output_image + + + +if __name__ == "__main__": + pass \ No newline at end of file diff --git a/custom_functions/image_text_to_image.py b/custom_functions/image_text_to_image.py new file mode 100644 index 000000000..5a828d996 --- /dev/null +++ b/custom_functions/image_text_to_image.py @@ -0,0 +1,156 @@ +import os +import random +import sys +from typing import Sequence, Mapping, Any, Union +import torch + + +def get_value_at_index(obj: Union[Sequence, Mapping], index: int) -> Any: + """Returns the value at the given index of a sequence or mapping. + + If the object is a sequence (like list or string), returns the value at the given index. + If the object is a mapping (like a dictionary), returns the value at the index-th key. + + Some return a dictionary, in these cases, we look for the "results" key + + Args: + obj (Union[Sequence, Mapping]): The object to retrieve the value from. + index (int): The index of the value to retrieve. + + Returns: + Any: The value at the given index. + + Raises: + IndexError: If the index is out of bounds for the object and the object is not a mapping. + """ + try: + return obj[index] + except KeyError: + return obj["result"][index] + + +def find_path(name: str, path: str = None) -> str: + """ + Recursively looks at parent folders starting from the given path until it finds the given name. + Returns the path as a Path object if found, or None otherwise. + """ + # If no path is given, use the current working directory + if path is None: + path = os.getcwd() + + # Check if the current directory contains the name + if name in os.listdir(path): + path_name = os.path.join(path, name) + print(f"{name} found: {path_name}") + return path_name + + # Get the parent directory + parent_directory = os.path.dirname(path) + + # If the parent directory is the same as the current directory, we've reached the root and stop the search + if parent_directory == path: + return None + + # Recursively call the function with the parent directory + return find_path(name, parent_directory) + + +def add_comfyui_directory_to_sys_path() -> None: + """ + Add 'ComfyUI' to the sys.path + """ + comfyui_path = find_path("ComfyUI") + if comfyui_path is not None and os.path.isdir(comfyui_path): + sys.path.append(comfyui_path) + print(f"'{comfyui_path}' added to sys.path") + + +def add_extra_model_paths() -> None: + """ + Parse the optional extra_model_paths.yaml file and add the parsed paths to the sys.path. + """ + from main import load_extra_path_config + + extra_model_paths = find_path("extra_model_paths.yaml") + + if extra_model_paths is not None: + load_extra_path_config(extra_model_paths) + else: + print("Could not find the extra_model_paths config file.") + +add_comfyui_directory_to_sys_path() +add_extra_model_paths() + + +from nodes import ( + NODE_CLASS_MAPPINGS, + SaveImage, + CLIPTextEncode, + LoadImage, + CheckpointLoaderSimple, + KSampler, + VAEDecode, + VAEEncode, +) + + +def image_text_to_image(image_path, pos_text, neg_text=r'watermark, text'): + + with torch.inference_mode(): + checkpointloadersimple = CheckpointLoaderSimple() + checkpointloadersimple_14 = checkpointloadersimple.load_checkpoint( + ckpt_name="v1-5-pruned-emaonly.ckpt" + ) + + cliptextencode = CLIPTextEncode() + cliptextencode_6 = cliptextencode.encode( + text=pos_text, + clip=get_value_at_index(checkpointloadersimple_14, 1), + ) + + cliptextencode_7 = cliptextencode.encode( + text=neg_text, + clip=get_value_at_index(checkpointloadersimple_14, 1), + ) + + loadimage = LoadImage() + loadimage_10 = loadimage.load_image(image=image_path, abs_path=True) + + vaeencode = VAEEncode() + vaeencode_12 = vaeencode.encode( + pixels=get_value_at_index(loadimage_10, 0), + vae=get_value_at_index(checkpointloadersimple_14, 2), + ) + + ksampler = KSampler() + vaedecode = VAEDecode() + saveimage = SaveImage() + + ksampler_3 = ksampler.sample( + seed=random.randint(1, 2**64), + steps=20, + cfg=8, + sampler_name="dpmpp_2m", + scheduler="normal", + denoise=0.8700000000000001, + model=get_value_at_index(checkpointloadersimple_14, 0), + positive=get_value_at_index(cliptextencode_6, 0), + negative=get_value_at_index(cliptextencode_7, 0), + latent_image=get_value_at_index(vaeencode_12, 0), + ) + + vaedecode_8 = vaedecode.decode( + samples=get_value_at_index(ksampler_3, 0), + vae=get_value_at_index(checkpointloadersimple_14, 2), + ) + + output_image = get_value_at_index(vaedecode_8, 0) + # print("output", output_image.shape, torch.max(output_image), torch.min(output_image)) + + return output_image + +if __name__ == "__main__": + image_path = r'D:\VisualForge\ComfyUI\input\example.png' + pos_text = r'photograph of victorian woman with wings, sky clouds, meadow grass' + + out_image = image_text_to_image(image_path, pos_text) \ No newline at end of file diff --git a/nodes.py b/nodes.py index 513cd0c7f..4be25aab7 100644 --- a/nodes.py +++ b/nodes.py @@ -1544,8 +1544,11 @@ class LoadImage: RETURN_TYPES = ("IMAGE", "MASK") FUNCTION = "load_image" - def load_image(self, image): - image_path = folder_paths.get_annotated_filepath(image) + def load_image(self, image, abs_path=False): + if not abs_path: + image_path = folder_paths.get_annotated_filepath(image) + else: + image_path = image img = node_helpers.pillow(Image.open, image_path)