From 9371924e654128258cc82419e83c2a788a32e2be Mon Sep 17 00:00:00 2001
From: missionfloyd <missionfloyd@users.noreply.github.com>
Date: Thu, 13 Apr 2023 03:11:17 -0600
Subject: [PATCH] Move mask conversion to separate file

---
 comfy_extras/nodes_mask_conversion.py | 54 +++++++++++++++++++++++++++
 nodes.py                              | 42 +--------------------
 2 files changed, 55 insertions(+), 41 deletions(-)
 create mode 100644 comfy_extras/nodes_mask_conversion.py

diff --git a/comfy_extras/nodes_mask_conversion.py b/comfy_extras/nodes_mask_conversion.py
new file mode 100644
index 00000000..04dcbd0d
--- /dev/null
+++ b/comfy_extras/nodes_mask_conversion.py
@@ -0,0 +1,54 @@
+import numpy as np
+import torch
+import torch.nn.functional as F
+from PIL import Image
+
+import comfy.utils
+
+class ImageToMask:
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+                "required": {
+                    "image": ("IMAGE",),
+                    "channel": (["red", "green", "blue"],),
+                }
+        }
+
+    CATEGORY = "image"
+
+    RETURN_TYPES = ("MASK",)
+    FUNCTION = "image_to_mask"
+
+    def image_to_mask(self, image, channel):
+        channels = ["red", "green", "blue"]
+        mask = image[0, :, :, channels.index(channel)]
+        return (mask,)
+
+class MaskToImage:
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+                "required": {
+                    "mask": ("MASK",),
+                }
+        }
+
+    CATEGORY = "image"
+
+    RETURN_TYPES = ("IMAGE",)
+    FUNCTION = "mask_to_image"
+
+    def mask_to_image(self, mask):
+        result = mask[None, :, :, None].expand(-1, -1, -1, 3)
+        return (result,)
+        
+NODE_CLASS_MAPPINGS = {
+    "ImageToMask": ImageToMask,
+    "MaskToImage": MaskToImage,
+}
+
+NODE_DISPLAY_NAME_MAPPINGS = {
+    "ImageToMask": "Convert Image to Mask",
+    "MaskToImage": "Convert Mask to Image",
+}
diff --git a/nodes.py b/nodes.py
index 325e3ba6..3ed9cf49 100644
--- a/nodes.py
+++ b/nodes.py
@@ -1061,43 +1061,6 @@ class ImagePadForOutpaint:
 
         return (new_image, mask)
 
-class ImageToMask:
-    @classmethod
-    def INPUT_TYPES(s):
-        return {
-                "required": {
-                    "image": ("IMAGE",),
-                    "channel": (["red", "green", "blue"],),
-                }
-        }
-
-    CATEGORY = "image"
-
-    RETURN_TYPES = ("MASK",)
-    FUNCTION = "image_to_mask"
-
-    def image_to_mask(self, image, channel):
-        channels = ["red", "green", "blue"]
-        mask = image[0, :, :, channels.index(channel)]
-        return (mask,)
-
-class MaskToImage:
-    @classmethod
-    def INPUT_TYPES(s):
-        return {
-                "required": {
-                    "mask": ("MASK",),
-                }
-        }
-
-    CATEGORY = "image"
-
-    RETURN_TYPES = ("IMAGE",)
-    FUNCTION = "mask_to_image"
-
-    def mask_to_image(self, mask):
-        result = mask[None, :, :, None].expand(-1, -1, -1, 3)
-        return (result,)
 
 NODE_CLASS_MAPPINGS = {
     "KSampler": KSampler,
@@ -1141,8 +1104,6 @@ NODE_CLASS_MAPPINGS = {
     "unCLIPCheckpointLoader": unCLIPCheckpointLoader,
     "CheckpointLoader": CheckpointLoader,
     "DiffusersLoader": DiffusersLoader,
-    "ImageToMask": ImageToMask,
-    "MaskToImage": MaskToImage,
 }
 
 NODE_DISPLAY_NAME_MAPPINGS = {
@@ -1188,8 +1149,6 @@ NODE_DISPLAY_NAME_MAPPINGS = {
     "ImageUpscaleWithModel": "Upscale Image (using Model)",
     "ImageInvert": "Invert Image",
     "ImagePadForOutpaint": "Pad Image for Outpainting",
-    "ImageToMask": "Convert Image to Mask",
-    "MaskToImage": "Convert Mask to Image",
     # _for_testing
     "VAEDecodeTiled": "VAE Decode (Tiled)",
     "VAEEncodeTiled": "VAE Encode (Tiled)",
@@ -1233,3 +1192,4 @@ def init_custom_nodes():
     load_custom_nodes()
     load_custom_node(os.path.join(os.path.join(os.path.dirname(os.path.realpath(__file__)), "comfy_extras"), "nodes_upscale_model.py"))
     load_custom_node(os.path.join(os.path.join(os.path.dirname(os.path.realpath(__file__)), "comfy_extras"), "nodes_post_processing.py"))
+    load_custom_node(os.path.join(os.path.join(os.path.dirname(os.path.realpath(__file__)), "comfy_extras"), "nodes_mask_conversion.py"))