diff --git a/py/defs/__init__.py b/py/defs/__init__.py index e69de29..d472097 100644 --- a/py/defs/__init__.py +++ b/py/defs/__init__.py @@ -0,0 +1,17 @@ +import glob +import importlib +import os + +from .captures import CAPTURE_FIELD_LIST +from .samplers import SAMPLERS + +dir_name = os.path.dirname(os.path.abspath(__file__)) +for module_path in glob.glob(dir_name + "/ext/*.py"): + module_name = os.path.basename(module_path) + module_name = os.path.splitext(module_name)[0] + package_name = ( + f"custom_nodes.ComfyUI-SaveImageWithMetaData.py.defs.ext.{module_name}" + ) + module = importlib.import_module(package_name) + CAPTURE_FIELD_LIST.update(getattr(module, "CAPTURE_FIELD_LIST", {})) + SAMPLERS.update(getattr(module, "SAMPLERS", {})) diff --git a/py/defs/ext/size_from_presets.py b/py/defs/ext/size_from_presets.py new file mode 100644 index 0000000..f667a5b --- /dev/null +++ b/py/defs/ext/size_from_presets.py @@ -0,0 +1,22 @@ +from ..meta import MetaField + + +def get_width(preset): + return preset.split("x")[0].strip() + + +def get_height(preset): + return preset.split("x")[1].strip() + + +CAPTURE_FIELD_LIST = { + "EmptyLatentImageFromPresetsSD15": { + MetaField.IMAGE_WIDTH: {"field_name": "preset", "format": get_width}, + MetaField.IMAGE_HEIGHT: {"field_name": "preset", "format": get_height}, + }, + "EmptyLatentImageFromPresetsSDXL": { + MetaField.IMAGE_WIDTH: {"field_name": "preset", "format": get_width}, + MetaField.IMAGE_HEIGHT: {"field_name": "preset", "format": get_height}, + }, + # TODO RandomEmptyLatentImageFromPresetsSD.. +} diff --git a/py/defs/meta.py b/py/defs/meta.py index 4f3ae56..250ee38 100644 --- a/py/defs/meta.py +++ b/py/defs/meta.py @@ -1,7 +1,7 @@ -from enum import Enum +from enum import IntEnum -class MetaField(Enum): +class MetaField(IntEnum): MODEL_NAME = 0 MODEL_HASH = 1 POSITIVE_PROMPT = 10 diff --git a/py/nodes/node.py b/py/nodes/node.py index 27b452d..be95122 100644 --- a/py/nodes/node.py +++ b/py/nodes/node.py @@ -57,11 +57,12 @@ class SaveImageWithMetaData(BaseNode): inputs = Capture.get_inputs() print("[get_inputs]") print(inputs) + + print("[PNGInfo Dict]") pnginfo_dict = Capture.gen_pnginfo_dict(inputs) if len(images) >= 2: pnginfo_dict["Batch index"] = index pnginfo_dict["Batch size"] = len(images) - print("[PNGInfo Dict]") print(pnginfo_dict) metadata = None