Fix category for all Core ML nodes

coreml_suite/__init__.py

@@ -0,0 +1,2 @@
+class COREML_NODE:
+    CATEGORY = "Core ML Suite"

coreml_suite/lcm/nodes.py

@@ -3,10 +3,11 @@ import os
 from coremltools import ComputeUnit
 from python_coreml_stable_diffusion.coreml_model import CoreMLModel
 
+from coreml_suite import COREML_NODE
 from coreml_suite.lcm import converter as lcm_converter
 
 
-class COREML_CONVERT_LCM:
+class COREML_CONVERT_LCM(COREML_NODE):
     """Converts a LCM model to Core ML."""
 
     @classmethod

coreml_suite/nodes.py

@@ -8,6 +8,7 @@ import folder_paths
 from comfy import model_base
 from comfy.model_management import get_torch_device
 from comfy.model_patcher import ModelPatcher
+from coreml_suite import COREML_NODE
 from coreml_suite.lcm.utils import add_lcm_model_options, lcm_patch, is_lcm
 from coreml_suite.logger import logger
 from nodes import KSampler
@@ -16,7 +17,7 @@ from coreml_suite.models import CoreMLModelWrapper
 from coreml_suite.config import get_model_config
 
 
-class CoreMLSampler(KSampler):
+class CoreMLSampler(COREML_NODE, KSampler):
     @classmethod
     def INPUT_TYPES(s):
         old_required = KSampler.INPUT_TYPES()["required"].copy()
@@ -29,8 +30,6 @@ class CoreMLSampler(KSampler):
             "optional": {"negative": ("CONDITIONING",), "latent_image": ("LATENT",)},
         }
 
-    CATEGORY = "Core ML Suite"
-
     def sample(
         self,
         coreml_model,
@@ -89,7 +88,7 @@ class CoreMLSampler(KSampler):
         return model_patcher
 
 
-class CoreMLLoader:
+class CoreMLLoader(COREML_NODE):
     PACKAGE_DIRNAME = ""
 
     @classmethod
@@ -109,7 +108,6 @@ class CoreMLLoader:
         }
 
     FUNCTION = "load"
-    CATEGORY = "Core ML Suite"
 
     @classmethod
     def coreml_filenames(cls):