57 lines
1.8 KiB
Python
57 lines
1.8 KiB
Python
from comfy.comfy_types import IO, ComfyNodeABC, InputTypeDict
|
|
import os
|
|
import folder_paths
|
|
import node_helpers
|
|
from PIL import Image, ImageOps, ImageSequence
|
|
import numpy as np
|
|
import torch
|
|
import comfy.model_management
|
|
|
|
# Import the original LoadImage node for delegation
|
|
try:
|
|
from nodes import LoadImage as OriginalLoadImage
|
|
except ImportError:
|
|
raise ImportError(
|
|
"Unable to import ComfyUI's LoadImage node. "
|
|
"This custom node requires ComfyUI to be installed and loaded first."
|
|
)
|
|
|
|
class LoadImageWithMetadata(ComfyNodeABC):
|
|
"""Load an image from disk while preserving ComfyUI workflow metadata from PNG.
|
|
|
|
This is a drop-in replacement for ComfyUI's built-in Load Image node,
|
|
but it preserves the embedded workflow metadata instead of stripping it.
|
|
"""
|
|
|
|
DESCRIPTION = """Load image from disk with ComfyUI workflow metadata preserved.
|
|
|
|
Works with PNG files that contain ComfyUI metadata.
|
|
"""
|
|
|
|
CATEGORY = "image/load"
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(s) -> InputTypeDict:
|
|
return OriginalLoadImage.INPUT_TYPES()
|
|
|
|
RETURN_TYPES = (*OriginalLoadImage.RETURN_TYPES, IO.STRING)
|
|
try:
|
|
RETURN_NAMES = (*OriginalLoadImage.RETURN_NAMES, "workflow_metadata")
|
|
except AttributeError:
|
|
RETURN_NAMES = (None,) # RETURN_NAMES appears to be optional
|
|
|
|
FUNCTION = "load_image"
|
|
OUTPUT_NODE = True
|
|
|
|
def load_image(self, image):
|
|
"""Load image and preserve metadata."""
|
|
|
|
# Call original implementation
|
|
ret = OriginalLoadImage.load_image(self, image)
|
|
|
|
# Load the original image to get metadata
|
|
image_path = folder_paths.get_annotated_filepath(image)
|
|
img = Image.open(image_path)
|
|
metadata = img.info.get("workflow", None)
|
|
|
|
return (*ret, metadata) |