feat: custom function support

ComfyUI-Manager

@@ -0,0 +1 @@
+Subproject commit d7170c0264ca1623b7a4ba7d63756a9b54a7d119

ComfyUI-to-Python-Extension

@@ -0,0 +1 @@
+Subproject commit d3b82726a7f185f25f3d0194702ea17df47ee167

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)

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

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)

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)