diff --git a/coreml_suite/__init__.py b/coreml_suite/__init__.py index e69de29..d652ee3 100644 --- a/coreml_suite/__init__.py +++ b/coreml_suite/__init__.py @@ -0,0 +1,2 @@ +class COREML_NODE: + CATEGORY = "Core ML Suite" diff --git a/coreml_suite/lcm/nodes.py b/coreml_suite/lcm/nodes.py index 2fc5b4b..6c7d885 100644 --- a/coreml_suite/lcm/nodes.py +++ b/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 diff --git a/coreml_suite/nodes.py b/coreml_suite/nodes.py index 80560e8..04b66a0 100644 --- a/coreml_suite/nodes.py +++ b/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):