Merge fdbcf4c02b into f3ea976cba

comfy/cli_args.py

@@ -91,6 +91,7 @@ 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,6 +1,8 @@
 import torch
 import logging
 
+from comfy.cli_args import args
+
 try:
     import comfy_kitchen as ck
     from comfy_kitchen.tensor import (
@@ -21,7 +23,15 @@ try:
             ck.registry.disable("cuda")
             logging.warning("WARNING: You need pytorch with cu130 or higher to use optimized CUDA operations.")
 
-    ck.registry.disable("triton")
+    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")
     for k, v in ck.list_backends().items():
         logging.info(f"Found comfy_kitchen backend {k}: {v}")
 except ImportError as e:

comfy_extras/nodes_images.py

@@ -192,7 +192,7 @@ class SaveAnimatedWEBP(IO.ComfyNode):
             category="image/animation",
             inputs=[
                 IO.Image.Input("images"),
-                IO.String.Input("filename_prefix", default="ComfyUI"),
+                IO.String.Input("filename_prefix", default="%date:yyyy-MM-dd%/ComfyUI"),
                 IO.Float.Input("fps", default=6.0, min=0.01, max=1000.0, step=0.01),
                 IO.Boolean.Input("lossless", default=True),
                 IO.Int.Input("quality", default=80, min=0, max=100),
@@ -229,7 +229,7 @@ class SaveAnimatedPNG(IO.ComfyNode):
             category="image/animation",
             inputs=[
                 IO.Image.Input("images"),
-                IO.String.Input("filename_prefix", default="ComfyUI"),
+                IO.String.Input("filename_prefix", default="%date:yyyy-MM-dd%/ComfyUI"),
                 IO.Float.Input("fps", default=6.0, min=0.01, max=1000.0, step=0.01),
                 IO.Int.Input("compress_level", default=4, min=0, max=9, advanced=True),
             ],

comfy_extras/nodes_post_processing.py

@@ -666,12 +666,13 @@ class ColorTransfer(io.ComfyNode):
     def define_schema(cls):
         return io.Schema(
             node_id="ColorTransfer",
+            display_name="Color Transfer",
             category="image/postprocessing",
             description="Match the colors of one image to another using various algorithms.",
             search_aliases=["color match", "color grading", "color correction", "match colors", "color transform", "mkl", "reinhard", "histogram"],
             inputs=[
                 io.Image.Input("image_target", tooltip="Image(s) to apply the color transform to."),
-                io.Image.Input("image_ref", optional=True, tooltip="Reference image(s) to match colors to. If not provided, processing is skipped"),
+                io.Image.Input("image_ref", tooltip="Reference image(s) to match colors to."),
                 io.Combo.Input("method", options=['reinhard_lab', 'mkl_lab', 'histogram'],),
                 io.DynamicCombo.Input("source_stats",
                     tooltip="per_frame: each frame matched to image_ref individually. uniform: pool stats across all source frames as baseline, match to image_ref. target_frame: use one chosen frame as the baseline for the transform to image_ref, applied uniformly to all frames (preserves relative differences)",

comfy_extras/nodes_primitive.py

@@ -49,7 +49,7 @@ class Int(io.ComfyNode):
             display_name="Int",
             category="utils/primitive",
             inputs=[
-                io.Int.Input("value", min=-sys.maxsize, max=sys.maxsize, control_after_generate=True),
+                io.Int.Input("value", min=-sys.maxsize, max=sys.maxsize, control_after_generate=io.ControlAfterGenerate.fixed),
             ],
             outputs=[io.Int.Output()],
         )

extra_model_paths.yaml.example

@@ -28,7 +28,7 @@
 #config for a1111 ui
 #all you have to do is uncomment this (remove the #) and change the base_path to where yours is installed
 
-#a111:
+#a1111:
 #     base_path: path/to/stable-diffusion-webui/
 #     checkpoints: models/Stable-diffusion
 #     configs: models/Stable-diffusion

folder_paths.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import os
+import re
 import time
 import mimetypes
 import logging
@@ -433,45 +434,100 @@ def get_save_image_path(filename_prefix: str, output_dir: str, image_width=0, im
         prefix = filename[:prefix_len + 1]
         try:
             digits = int(filename[prefix_len + 1:].split('_')[0])
-        except:
+        except ValueError:
             digits = 0
         return digits, prefix
 
-    def compute_vars(input: str, image_width: int, image_height: int) -> str:
-        input = input.replace("%width%", str(image_width))
-        input = input.replace("%height%", str(image_height))
-        now = time.localtime()
-        input = input.replace("%year%", str(now.tm_year))
-        input = input.replace("%month%", str(now.tm_mon).zfill(2))
-        input = input.replace("%day%", str(now.tm_mday).zfill(2))
-        input = input.replace("%hour%", str(now.tm_hour).zfill(2))
-        input = input.replace("%minute%", str(now.tm_min).zfill(2))
-        input = input.replace("%second%", str(now.tm_sec).zfill(2))
-        return input
+    def compute_vars(input_str: str, image_width: int, image_height: int) -> str:
+        date_pattern = re.compile(r'%date:(.+?)%')
 
-    if "%" in filename_prefix:
-        filename_prefix = compute_vars(filename_prefix, image_width, image_height)
+        def replace_date(match):
+            date_format = match.group(1)
+
+            dateformat_conversion_map = {
+                "yyyy": "%Y",
+                "yy": "%y",
+                "MM": "%m",
+                "dd": "%d",
+                "hh": "%H",
+                "mm": "%M",
+                "ss": "%S",
+            }
+
+            strftime_format = date_format
+            for old, new in dateformat_conversion_map.items():
+                strftime_format = strftime_format.replace(old, new)
+
+            now = time.localtime()
+            formatted_date = time.strftime(strftime_format, now)
+            return formatted_date
+
+        input_str = date_pattern.sub(replace_date, input_str)
+
+        now = time.localtime()
+        input_str = input_str.replace("%year%", str(now.tm_year))
+        input_str = input_str.replace("%month%", str(now.tm_mon).zfill(2))
+        input_str = input_str.replace("%day%", str(now.tm_mday).zfill(2))
+        input_str = input_str.replace("%hour%", str(now.tm_hour).zfill(2))
+        input_str = input_str.replace("%minute%", str(now.tm_min).zfill(2))
+        input_str = input_str.replace("%second%", str(now.tm_sec).zfill(2))
+
+        input_str = input_str.replace("%width%", str(image_width))
+        input_str = input_str.replace("%height%", str(image_height))
+
+        return input_str
+
+    filename_prefix = compute_vars(filename_prefix, image_width, image_height)
 
     subfolder = os.path.dirname(os.path.normpath(filename_prefix))
     filename = os.path.basename(os.path.normpath(filename_prefix))
 
+    if not subfolder and ("%date:" in filename_prefix or "%year%" in filename_prefix or
+                         "%month%" in filename_prefix or "%day%" in filename_prefix or
+                         "%hour%" in filename_prefix or "%minute%" in filename_prefix or
+                         "%second%" in filename_prefix or "ComfyUI" in filename_prefix):
+        import locale
+        try:
+            locale.setlocale(locale.LC_TIME, '')
+            default_date_format = locale.nl_langinfo(locale.D_FMT).replace("%y", "%Y")
+            if os.name == 'nt':
+                default_date_format = default_date_format.replace("/", "-").replace("\\", "-").replace(":","-")
+            subfolder = time.strftime(default_date_format)
+
+        except (locale.Error, AttributeError):
+            subfolder = time.strftime("%Y-%m-%d")
+            if os.name == 'nt':
+                subfolder = subfolder.replace("/", "-").replace("\\", "-").replace(":","-")
+
     full_output_folder = os.path.join(output_dir, subfolder)
 
     if os.path.commonpath((output_dir, os.path.abspath(full_output_folder))) != output_dir:
-        err = "**** ERROR: Saving image outside the output folder is not allowed." + \
-              "\n full_output_folder: " + os.path.abspath(full_output_folder) + \
-              "\n         output_dir: " + output_dir + \
-              "\n         commonpath: " + os.path.commonpath((output_dir, os.path.abspath(full_output_folder)))
+        err = ("**** ERROR: Saving outside the output folder is not allowed."
+               f"\n full_output_folder: {os.path.abspath(full_output_folder)}"
+               f"\n         output_dir: {output_dir}"
+               f"\n         commonpath: {os.path.commonpath((output_dir, os.path.abspath(full_output_folder)))}")
         logging.error(err)
         raise Exception(err)
 
     try:
-        counter = max(filter(lambda a: os.path.normcase(a[1][:-1]) == os.path.normcase(filename) and a[1][-1] == "_", map(map_filename, os.listdir(full_output_folder))))[0] + 1
-    except ValueError:
-        counter = 1
+        files = [f for f in os.listdir(full_output_folder) if os.path.isfile(os.path.join(full_output_folder, f))]
+        mapped_files = [map_filename(f) for f in files]
+        filtered_files = [
+            (digits, prefix) for digits, prefix in mapped_files
+            if os.path.normcase(prefix[:-1]) == os.path.normcase(filename) and prefix[-1] == "_"
+        ]
+
+        if filtered_files:
+            counter = max(digits for digits, _ in filtered_files) + 1
+        else:
+            counter = 1
     except FileNotFoundError:
         os.makedirs(full_output_folder, exist_ok=True)
         counter = 1
+    except OSError as e:
+        logging.error(f"Error accessing directory: {e}")
+        raise
+
     return full_output_folder, filename, counter, subfolder, filename_prefix
 
 def get_input_subfolders() -> list[str]:

nodes.py

@@ -1617,7 +1617,7 @@ class SaveImage:
         return {
             "required": {
                 "images": ("IMAGE", {"tooltip": "The images to save."}),
-                "filename_prefix": ("STRING", {"default": "ComfyUI", "tooltip": "The prefix for the file to save. This may include formatting information such as %date:yyyy-MM-dd% or %Empty Latent Image.width% to include values from nodes."})
+                "filename_prefix": ("STRING", {"default": "%date:yyyy-MM-dd%/ComfyUI", "tooltip": "The prefix for the file to save. This may include formatting information such as '%date:yyyy-MM-dd%', '%year%-%month%-%day%' or %Empty Latent Image.width% to include values from nodes."})
             },
             "hidden": {
                 "prompt": "PROMPT", "extra_pnginfo": "EXTRA_PNGINFO"
@@ -1754,57 +1754,49 @@ class LoadImage:
 
         return True
 
-class LoadImageMask:
+
+class LoadImageMask(LoadImage):
     ESSENTIALS_CATEGORY = "Image Tools"
     SEARCH_ALIASES = ["import mask", "alpha mask", "channel mask"]
 
     _color_channels = ["alpha", "red", "green", "blue"]
+
     @classmethod
     def INPUT_TYPES(s):
-        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, ), }
-                }
+        types = super().INPUT_TYPES()
+        return {
+            "required": {
+                **types["required"],
+                "channel": (s._color_channels, )
+            }
+        }
 
     CATEGORY = "mask"
-
     RETURN_TYPES = ("MASK",)
-    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
+    FUNCTION = "load_image_mask"
+
+    def load_image_mask(self, image, channel):
+        image_tensor, mask_tensor = super().load_image(image)
         c = channel[0].upper()
-        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
+
+        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(),)
         else:
-            mask = torch.zeros((64,64), dtype=torch.float32, device="cpu")
-        return (mask.unsqueeze(0),)
+            empty_mask = torch.zeros(
+                image_tensor.shape[:-1],
+                dtype=image_tensor.dtype,
+                device=image_tensor.device
+            )
+            return (empty_mask,)
 
     @classmethod
     def IS_CHANGED(s, image, channel):
-        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
+        return super().IS_CHANGED(image)
 
 
 class LoadImageOutput(LoadImage):