comfy/cli_args.py

@@ -91,7 +91,6 @@ parser.add_argument("--directml", type=int, nargs="?", metavar="DIRECTML_DEVICE"
 
 parser.add_argument("--oneapi-device-selector", type=str, default=None, metavar="SELECTOR_STRING", help="Sets the oneAPI device(s) this instance will use.")
 parser.add_argument("--supports-fp8-compute", action="store_true", help="ComfyUI will act like if the device supports fp8 compute.")
-parser.add_argument("--enable-triton-backend", action="store_true", help="ComfyUI will enable the use of Triton backend in comfy-kitchen. Is disabled at launch by default.")
 
 class LatentPreviewMethod(enum.Enum):
     NoPreviews = "none"

comfy/quant_ops.py

@@ -1,8 +1,6 @@
 import torch
 import logging
 
-from comfy.cli_args import args
-
 try:
     import comfy_kitchen as ck
     from comfy_kitchen.tensor import (
@@ -23,15 +21,7 @@ try:
             ck.registry.disable("cuda")
             logging.warning("WARNING: You need pytorch with cu130 or higher to use optimized CUDA operations.")
 
-    if args.enable_triton_backend:
-        try:
-            import triton
-            logging.info("Found triton %s. Enabling comfy-kitchen triton backend.", triton.__version__)
-        except ImportError as e:
-            logging.error(f"Failed to import triton, Error: {e}, the comfy-kitchen triton backend will not be available.")
-            ck.registry.disable("triton")
-    else:
-        ck.registry.disable("triton")
+    ck.registry.disable("triton")
     for k, v in ck.list_backends().items():
         logging.info(f"Found comfy_kitchen backend {k}: {v}")
 except ImportError as e:

nodes.py

@@ -1754,49 +1754,57 @@ class LoadImage:
 
         return True
 
-
-class LoadImageMask(LoadImage):
+class LoadImageMask:
     ESSENTIALS_CATEGORY = "Image Tools"
     SEARCH_ALIASES = ["import mask", "alpha mask", "channel mask"]
 
     _color_channels = ["alpha", "red", "green", "blue"]
-
     @classmethod
     def INPUT_TYPES(s):
-        types = super().INPUT_TYPES()
-        return {
-            "required": {
-                **types["required"],
-                "channel": (s._color_channels, )
-            }
-        }
+        input_dir = folder_paths.get_input_directory()
+        files = [f for f in os.listdir(input_dir) if os.path.isfile(os.path.join(input_dir, f))]
+        return {"required":
+                    {"image": (sorted(files), {"image_upload": True}),
+                     "channel": (s._color_channels, ), }
+                }
 
     CATEGORY = "mask"
+
     RETURN_TYPES = ("MASK",)
-    FUNCTION = "load_image_mask"
-
-    def load_image_mask(self, image, channel):
-        image_tensor, mask_tensor = super().load_image(image)
+    FUNCTION = "load_image"
+    def load_image(self, image, channel):
+        image_path = folder_paths.get_annotated_filepath(image)
+        i = node_helpers.pillow(Image.open, image_path)
+        i = node_helpers.pillow(ImageOps.exif_transpose, i)
+        if i.getbands() != ("R", "G", "B", "A"):
+            if i.mode == 'I':
+                i = i.point(lambda i: i * (1 / 255))
+            i = i.convert("RGBA")
+        mask = None
         c = channel[0].upper()
-
-        if c == 'A':
-            return (mask_tensor,)
-
-        channel_idx = {'R': 0, 'G': 1, 'B': 2}.get(c, 0)
-
-        if channel_idx < image_tensor.shape[-1]:
-            return (image_tensor[..., channel_idx].clone(),)
+        if c in i.getbands():
+            mask = np.array(i.getchannel(c)).astype(np.float32) / 255.0
+            mask = torch.from_numpy(mask)
+            if c == 'A':
+                mask = 1. - mask
         else:
-            empty_mask = torch.zeros(
-                image_tensor.shape[:-1],
-                dtype=image_tensor.dtype,
-                device=image_tensor.device
-            )
-            return (empty_mask,)
+            mask = torch.zeros((64,64), dtype=torch.float32, device="cpu")
+        return (mask.unsqueeze(0),)
 
     @classmethod
     def IS_CHANGED(s, image, channel):
-        return super().IS_CHANGED(image)
+        image_path = folder_paths.get_annotated_filepath(image)
+        m = hashlib.sha256()
+        with open(image_path, 'rb') as f:
+            m.update(f.read())
+        return m.digest().hex()
+
+    @classmethod
+    def VALIDATE_INPUTS(s, image):
+        if not folder_paths.exists_annotated_filepath(image):
+            return "Invalid image file: {}".format(image)
+
+        return True
 
 
 class LoadImageOutput(LoadImage):