687 lines
20 KiB
Python
687 lines
20 KiB
Python
import copy
|
|
import math
|
|
import os
|
|
import random
|
|
import sys
|
|
import traceback
|
|
import shlex
|
|
import numpy as np
|
|
import torch
|
|
import hashlib
|
|
import json
|
|
|
|
from PIL import Image, ImageOps
|
|
from comfy.cli_args import args
|
|
from PIL.PngImagePlugin import PngInfo
|
|
from datetime import datetime
|
|
import pickle
|
|
import re
|
|
import folder_paths
|
|
|
|
def clean_filename(filename: str) -> str:
|
|
# Remove any character that is not a-z, A-Z, 0-9, space, dot, underscore, or hyphen
|
|
cleaned = re.sub(r'[^\w\s.-]', '', filename)
|
|
# Optionally replace spaces with underscores (or remove)
|
|
cleaned = cleaned.replace(' ', '_')
|
|
return cleaned
|
|
|
|
def pil2tensor(image):
|
|
return torch.from_numpy(np.array(image).astype(np.float32) / 255.0)
|
|
|
|
def ensure_stack_dir(stackpath, key):
|
|
path = os.path.join(stackpath, clean_filename(key))
|
|
os.makedirs(path, exist_ok=True)
|
|
return path
|
|
|
|
class AnyType(str):
|
|
def __ne__(self, __value: object) -> bool:
|
|
return False
|
|
|
|
any = AnyType("*")
|
|
|
|
class StackPushImage:
|
|
|
|
def __init__(self):
|
|
pass
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"key": ("STRING", {"default": "default"}),
|
|
"input_image": ("IMAGE",)
|
|
},
|
|
"optional": {
|
|
"stackpath": ("STRING",{"default": "./ComfyUI/Stack/"}),
|
|
"input_image": ("IMAGE",)
|
|
},
|
|
"hidden": {
|
|
"prompt": "PROMPT", "extra_pnginfo": "EXTRA_PNGINFO"
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ()
|
|
|
|
FUNCTION = "execute"
|
|
|
|
CATEGORY = "ConCarne"
|
|
|
|
OUTPUT_NODE = True
|
|
|
|
@classmethod
|
|
def IS_CHANGED(cls, **kwargs):
|
|
return float("NaN")
|
|
|
|
def execute(self, key, input_image, stackpath="./ComfyUI/Stack/", prompt=None, extra_pnginfo=None):
|
|
if (key == ""):
|
|
key = "default"
|
|
|
|
path = ensure_stack_dir(stackpath,key)
|
|
|
|
counter = 0
|
|
|
|
for (batch_number, image) in enumerate(input_image):
|
|
now = datetime.now()
|
|
|
|
i = 255. * image.cpu().numpy()
|
|
img = Image.fromarray(np.clip(i, 0, 255).astype(np.uint8))
|
|
metadata = None
|
|
if not args.disable_metadata:
|
|
metadata = PngInfo()
|
|
if prompt is not None:
|
|
metadata.add_text("prompt", json.dumps(prompt))
|
|
|
|
filename = "stack" + now.strftime("%Y%m%d%H%M%S%f") + str(batch_number) + ".png"
|
|
img.save(os.path.join(path, filename), pnginfo=metadata, compress_level=8)
|
|
counter += 1
|
|
|
|
return ( None, )
|
|
|
|
class StackPushString:
|
|
|
|
def __init__(self):
|
|
pass
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"key": ("STRING", {"default": "default"}),
|
|
"stringvalue": ("STRING",)
|
|
},
|
|
"optional": {
|
|
"stackpath": ("STRING",{"default": "./ComfyUI/Stack/"}),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ()
|
|
|
|
FUNCTION = "execute"
|
|
|
|
CATEGORY = "ConCarne"
|
|
|
|
OUTPUT_NODE = True
|
|
|
|
@classmethod
|
|
def IS_CHANGED(cls, **kwargs):
|
|
return float("NaN")
|
|
|
|
def execute(self, key, stringvalue, stackpath="./ComfyUI/Stack/"):
|
|
if (key == ""):
|
|
key = "default"
|
|
|
|
path = ensure_stack_dir(stackpath,key)
|
|
|
|
#for (batch_number, stringv) in enumerate(stringvalue):
|
|
now = datetime.now()
|
|
|
|
filename = os.path.join(path,"stack" + now.strftime("%Y%m%d%H%M%S%f") + "._st")
|
|
with open(filename, 'w', encoding='utf-8') as f:
|
|
f.write(stringvalue)
|
|
|
|
return ( None, )
|
|
|
|
class StackPushInt:
|
|
|
|
def __init__(self):
|
|
pass
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"key": ("STRING", {"default": "default"}),
|
|
"numbervalue": ("INT",)
|
|
},
|
|
"optional": {
|
|
"stackpath": ("STRING",{"default": "./ComfyUI/Stack/"}),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ()
|
|
|
|
FUNCTION = "execute"
|
|
|
|
CATEGORY = "ConCarne"
|
|
|
|
OUTPUT_NODE = True
|
|
|
|
@classmethod
|
|
def IS_CHANGED(cls, **kwargs):
|
|
return float("NaN")
|
|
|
|
def execute(self, key, numbervalue, stackpath="./ComfyUI/Stack/"):
|
|
if (key == ""):
|
|
key = "default"
|
|
|
|
path = ensure_stack_dir(stackpath,key)
|
|
|
|
# for (batch_number, num) in enumerate(numbervalue):
|
|
now = datetime.now()
|
|
filename = os.path.join(path,"stack" + now.strftime("%Y%m%d%H%M%S%f") + "._nu")
|
|
with open(filename, 'w', encoding='utf-8') as f:
|
|
f.write(str(numbervalue))
|
|
|
|
return ( None, )
|
|
|
|
class StackPushFloat:
|
|
|
|
def __init__(self):
|
|
pass
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"key": ("STRING", {"default": "default"}),
|
|
"numbervalue": ("FLOAT",)
|
|
},
|
|
"optional": {
|
|
"stackpath": ("STRING",{"default": "./ComfyUI/Stack/"}),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ()
|
|
|
|
FUNCTION = "execute"
|
|
|
|
CATEGORY = "ConCarne"
|
|
|
|
OUTPUT_NODE = True
|
|
|
|
@classmethod
|
|
def IS_CHANGED(cls, **kwargs):
|
|
return float("NaN")
|
|
|
|
def execute(self, key, numbervalue, stackpath="./ComfyUI/Stack/"):
|
|
if (key == ""):
|
|
key = "default"
|
|
|
|
path = ensure_stack_dir(stackpath,key)
|
|
|
|
# for (batch_number, num) in enumerate(numbervalue):
|
|
now = datetime.now()
|
|
filename = os.path.join(path,"stack" + now.strftime("%Y%m%d%H%M%S%f") + "._fl")
|
|
with open(filename, 'w', encoding='utf-8') as f:
|
|
f.write(str(numbervalue))
|
|
|
|
return ( None, )
|
|
|
|
class StackPushObject:
|
|
|
|
def __init__(self):
|
|
pass
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"key": ("STRING", {"default": "default"}),
|
|
"object": (any, {})
|
|
},
|
|
"optional": {
|
|
"stackpath": ("STRING",{"default": "./ComfyUI/Stack/"}),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ()
|
|
|
|
FUNCTION = "execute"
|
|
|
|
CATEGORY = "ConCarne"
|
|
|
|
OUTPUT_NODE = True
|
|
|
|
@classmethod
|
|
def IS_CHANGED(cls, **kwargs):
|
|
return float("NaN")
|
|
|
|
def execute(self, key, object, stackpath="./ComfyUI/Stack/"):
|
|
if (key == ""):
|
|
key = "default"
|
|
|
|
path = ensure_stack_dir(stackpath,key)
|
|
|
|
# for (batch_number, num) in enumerate(numbervalue):
|
|
now = datetime.now()
|
|
filename = os.path.join(path,"stack" + now.strftime("%Y%m%d%H%M%S%f") + "._pkl")
|
|
|
|
with open(filename, 'wb') as f:
|
|
pickle.dump(object, f)
|
|
|
|
return ( None, )
|
|
|
|
|
|
|
|
class StackPopImage:
|
|
|
|
def __init__(self):
|
|
pass
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"key": ("STRING", {"default": "default"}),
|
|
"fallback_enabled": ("BOOLEAN", {"default": False}),
|
|
"override_enabled": ("BOOLEAN", {"default": False}),
|
|
"stackmode": ( ['Normal', 'Peek', 'Pop Bottom', 'Peek Bottom'], {"default": "Normal"} ),
|
|
},
|
|
|
|
"optional": {
|
|
"stackpath": ("STRING",{"default": "./ComfyUI/Stack/"}),
|
|
"fallback_item": ("IMAGE",),
|
|
"override_item": ("IMAGE",),
|
|
},
|
|
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE",)
|
|
|
|
FUNCTION = "execute"
|
|
|
|
CATEGORY = "ConCarne"
|
|
|
|
@classmethod
|
|
def IS_CHANGED(cls, **kwargs):
|
|
return float("NaN")
|
|
|
|
def execute(self, key, fallback_enabled, override_enabled, stackmode, fallback_item=None, override_item=None, stackpath="./ComfyUI/Stack/"):
|
|
|
|
if override_enabled:
|
|
return (override_item,)
|
|
|
|
path = os.path.join(stackpath,key)
|
|
|
|
images = []
|
|
|
|
if not os.path.exists(path):
|
|
if (fallback_enabled):
|
|
return ( fallback_item, )
|
|
else:
|
|
raise Exception(f"Key doesn't exist")
|
|
|
|
# Get all .png files sorted by name
|
|
sorted_files = sorted([f for f in os.listdir(path) if f.lower().endswith('.png')] )
|
|
|
|
last_file = ""
|
|
|
|
# Check if there are any PNG files
|
|
if sorted_files:
|
|
if stackmode == 'Normal' or stackmode == 'Peek':
|
|
last_file = os.path.join(path, sorted_files[-1])
|
|
else:
|
|
last_file = os.path.join(path, sorted_files[0])
|
|
|
|
image = Image.open(last_file)
|
|
#print(f"Opened: {last_file}")
|
|
else:
|
|
if (fallback_enabled):
|
|
return ( fallback_item, )
|
|
else:
|
|
raise Exception("No image files found. Stack Empty.")
|
|
|
|
image = ImageOps.exif_transpose(image)
|
|
image = image.convert("RGB")
|
|
image = pil2tensor(image).unsqueeze(0)
|
|
|
|
images.append(image)
|
|
|
|
if stackmode == 'Normal' or stackmode == 'Pop Bottom':
|
|
os.remove(last_file)
|
|
|
|
return ( torch.cat(images, dim=0), )
|
|
|
|
class StackPopString:
|
|
|
|
def __init__(self):
|
|
pass
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"key": ("STRING", {"default": "default"}),
|
|
"fallback_enabled": ("BOOLEAN", {"default": False}),
|
|
"override_enabled": ("BOOLEAN", {"default": False}),
|
|
"stackmode": ( ['Normal', 'Peek', 'Pop Bottom', 'Peek Bottom'], {"default": "Normal"} ),
|
|
},
|
|
"optional": {
|
|
"stackpath": ("STRING",{"default": "./ComfyUI/Stack/"}),
|
|
"fallback_item": ("STRING",),
|
|
"override_item": ("STRING",),
|
|
},
|
|
|
|
}
|
|
|
|
RETURN_TYPES = ("STRING",)
|
|
|
|
FUNCTION = "execute"
|
|
|
|
CATEGORY = "ConCarne"
|
|
|
|
@classmethod
|
|
def IS_CHANGED(cls, **kwargs):
|
|
return float("NaN")
|
|
|
|
def execute(self, key, fallback_enabled, fallback_item, override_enabled, override_item, stackmode, stackpath="./ComfyUI/Stack/"):
|
|
|
|
if override_enabled:
|
|
return (override_item,)
|
|
|
|
path = os.path.join(stackpath,key)
|
|
|
|
images = []
|
|
|
|
if not os.path.exists(path):
|
|
if (fallback_enabled):
|
|
return ( fallback_item, )
|
|
else:
|
|
raise Exception(f"Key doesn't exist")
|
|
|
|
# Get all ._st files sorted by name
|
|
sorted_files = sorted([f for f in os.listdir(path) if f.lower().endswith('._st')] )
|
|
|
|
last_file = ""
|
|
|
|
result = ""
|
|
|
|
# Check if there are any string files
|
|
if sorted_files:
|
|
if stackmode == 'Normal' or stackmode == 'Peek':
|
|
last_file = os.path.join(path, sorted_files[-1])
|
|
else:
|
|
last_file = os.path.join(path, sorted_files[0])
|
|
|
|
with open(last_file, 'r', encoding='utf-8') as f:
|
|
result = f.read()
|
|
#print(f"Opened: {last_file}")
|
|
else:
|
|
if (fallback_enabled):
|
|
return ( fallback_item, )
|
|
else:
|
|
raise Exception("No string files found. Stack Empty.")
|
|
|
|
if stackmode == 'Normal' or stackmode == 'Pop Bottom':
|
|
os.remove(last_file)
|
|
|
|
return (result, )
|
|
|
|
class StackPopInt:
|
|
|
|
def __init__(self):
|
|
pass
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"key": ("STRING", {"default": "default"}),
|
|
"fallback_enabled": ("BOOLEAN", {"default": False}),
|
|
"override_enabled": ("BOOLEAN", {"default": False}),
|
|
"stackmode": ( ['Normal', 'Peek', 'Pop Bottom', 'Peek Bottom'], {"default": "Normal"} ),
|
|
},
|
|
"optional": {
|
|
"stackpath": ("STRING",{"default": "./ComfyUI/Stack/"}),
|
|
"fallback_item": ("INT",),
|
|
"override_item": ("INT",),
|
|
},
|
|
|
|
}
|
|
|
|
RETURN_TYPES = ("INT",)
|
|
|
|
FUNCTION = "execute"
|
|
|
|
CATEGORY = "ConCarne"
|
|
|
|
@classmethod
|
|
def IS_CHANGED(cls, **kwargs):
|
|
return float("NaN")
|
|
|
|
def execute(self, key, fallback_enabled, fallback_item, override_enabled, override_item, stackmode, stackpath="./ComfyUI/Stack/"):
|
|
|
|
if override_enabled:
|
|
return (override_item,)
|
|
|
|
path = os.path.join(stackpath,key)
|
|
|
|
images = []
|
|
|
|
if not os.path.exists(path):
|
|
if (fallback_enabled):
|
|
return ( fallback_item, )
|
|
else:
|
|
raise Exception(f"Key doesn't exist")
|
|
|
|
# Get all ._nu files sorted by name
|
|
sorted_files = sorted([f for f in os.listdir(path) if f.lower().endswith('._nu')] )
|
|
|
|
last_file = ""
|
|
|
|
result = ""
|
|
|
|
# Check if there are any int files
|
|
if sorted_files:
|
|
if stackmode == 'Normal' or stackmode == 'Peek':
|
|
last_file = os.path.join(path, sorted_files[-1])
|
|
else:
|
|
last_file = os.path.join(path, sorted_files[0])
|
|
|
|
with open(last_file, 'r', encoding='utf-8') as f:
|
|
result = f.read()
|
|
#print(f"Opened: {last_file}")
|
|
else:
|
|
if (fallback_enabled):
|
|
return ( fallback_item, )
|
|
else:
|
|
raise Exception("No integer files found. Stack Empty.")
|
|
|
|
if stackmode == 'Normal' or stackmode == 'Pop Bottom':
|
|
os.remove(last_file)
|
|
|
|
return (int(result), )
|
|
|
|
class StackPopFloat:
|
|
|
|
def __init__(self):
|
|
pass
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"key": ("STRING", {"default": "default"}),
|
|
"fallback_enabled": ("BOOLEAN", {"default": False}),
|
|
"override_enabled": ("BOOLEAN", {"default": False}),
|
|
"stackmode": ( ['Normal', 'Peek', 'Pop Bottom', 'Peek Bottom'], {"default": "Normal"} ),
|
|
},
|
|
"optional": {
|
|
"stackpath": ("STRING",{"default": "./ComfyUI/Stack/"}),
|
|
"fallback_item": ("FLOAT",),
|
|
"override_item": ("FLOAT",),
|
|
},
|
|
|
|
}
|
|
|
|
RETURN_TYPES = ("FLOAT",)
|
|
|
|
FUNCTION = "execute"
|
|
|
|
CATEGORY = "ConCarne"
|
|
|
|
@classmethod
|
|
def IS_CHANGED(cls, **kwargs):
|
|
return float("NaN")
|
|
|
|
def execute(self, key, fallback_enabled, fallback_item, override_enabled, stackmode, override_item, stackpath="./ComfyUI/Stack/"):
|
|
|
|
if override_enabled:
|
|
return (override_item,)
|
|
|
|
path = os.path.join(stackpath,key)
|
|
|
|
images = []
|
|
|
|
if not os.path.exists(path):
|
|
if (fallback_enabled):
|
|
return ( fallback_item, )
|
|
else:
|
|
raise Exception(f"Key doesn't exist")
|
|
|
|
# Get all ._fl files sorted by name
|
|
sorted_files = sorted([f for f in os.listdir(path) if f.lower().endswith('._fl')] )
|
|
|
|
last_file = ""
|
|
|
|
result = ""
|
|
|
|
# Check if there are any float files
|
|
if sorted_files:
|
|
if stackmode == 'Normal' or stackmode == 'Peek':
|
|
last_file = os.path.join(path, sorted_files[-1])
|
|
else:
|
|
last_file = os.path.join(path, sorted_files[0])
|
|
with open(last_file, 'r', encoding='utf-8') as f:
|
|
result = f.read()
|
|
#print(f"Opened: {last_file}")
|
|
else:
|
|
if (fallback_enabled):
|
|
return ( fallback_item, )
|
|
else:
|
|
raise Exception("No float files found. Stack Empty.")
|
|
|
|
if stackmode == 'Normal' or stackmode == 'Pop Bottom':
|
|
os.remove(last_file)
|
|
|
|
return (float(result), )
|
|
|
|
class StackPopObject:
|
|
|
|
def __init__(self):
|
|
pass
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"key": ("STRING", {"default": "default"}),
|
|
"fallback_enabled": ("BOOLEAN", {"default": False}),
|
|
"override_enabled": ("BOOLEAN", {"default": False}),
|
|
"stackmode": ( ['Normal', 'Peek', 'Pop Bottom', 'Peek Bottom'], {"default": "Normal"} ),
|
|
},
|
|
"optional": {
|
|
"stackpath": ("STRING",{"default": "./ComfyUI/Stack/"}),
|
|
"fallback_item": (any, {}),
|
|
"override_item": (any, {}),
|
|
},
|
|
|
|
}
|
|
|
|
RETURN_TYPES = (any,)
|
|
RETURN_NAMES = ("any",)
|
|
# RETURN_TYPES = ("CONDITIONING","LATENT","IMAGE","STRING")
|
|
# OUTPUT_NODE = True
|
|
FUNCTION = "execute"
|
|
|
|
CATEGORY = "ConCarne"
|
|
|
|
@classmethod
|
|
def IS_CHANGED(cls, **kwargs):
|
|
return float("NaN")
|
|
|
|
def execute(self, key, fallback_enabled, override_enabled, stackmode, override_item=None, fallback_item=None, stackpath="./ComfyUI/Stack/"):
|
|
|
|
if not os.path.exists(os.path.join(folder_paths.get_user_directory(), "allow_generic_types.plz")):
|
|
raise Exception("Generic object popping not allowed until dummy file \"allow_generic_types.plz\" exists in user folder ("+folder_paths.get_user_directory()+"). You should not do this unless you understand the risks of loading generic objects.")
|
|
|
|
if override_enabled:
|
|
return (override_item,)
|
|
|
|
path = os.path.join(stackpath,key)
|
|
|
|
images = []
|
|
|
|
if not os.path.exists(path):
|
|
if (fallback_enabled):
|
|
return ( fallback_item, )
|
|
else:
|
|
raise Exception(f"Key doesn't exist")
|
|
|
|
# Get all .pkl files sorted by name
|
|
sorted_files = sorted([f for f in os.listdir(path) if f.lower().endswith('._pkl')] )
|
|
|
|
last_file = ""
|
|
|
|
result = ""
|
|
|
|
# Check if there are any pickle files
|
|
if sorted_files:
|
|
if stackmode == 'Normal' or stackmode == 'Peek':
|
|
last_file = os.path.join(path, sorted_files[-1])
|
|
else:
|
|
last_file = os.path.join(path, sorted_files[0])
|
|
|
|
with open(last_file, 'rb') as f:
|
|
loaded_object = pickle.load(f)
|
|
#print(f"Opened: {last_file}")
|
|
else:
|
|
if (fallback_enabled):
|
|
return ( fallback_item, )
|
|
else:
|
|
raise Exception("No object files found. Stack Empty.")
|
|
|
|
if stackmode == 'Normal' or stackmode == 'Pop Bottom':
|
|
os.remove(last_file)
|
|
|
|
return (loaded_object,)#loaded_object,loaded_object,loaded_object )
|
|
|
|
# A dictionary that contains all nodes you want to export with their names
|
|
# NOTE: names should be globally unique
|
|
NODE_CLASS_MAPPINGS = {
|
|
"StackPushImage": StackPushImage,
|
|
"StackPushString": StackPushString,
|
|
"StackPushInt": StackPushInt,
|
|
"StackPushFloat": StackPushFloat,
|
|
"StackPushObject": StackPushObject,
|
|
"StackPopImage": StackPopImage,
|
|
"StackPopString": StackPopString,
|
|
"StackPopInt": StackPopInt,
|
|
"StackPopFloat": StackPopFloat,
|
|
"StackPopObject": StackPopObject,
|
|
}
|
|
|
|
# A dictionary that contains the friendly/humanly readable titles for the nodes
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"StackPushImage": "StackPushImage",
|
|
"StackPushString": "StackPushString",
|
|
"StackPushInt": "StackPushInt",
|
|
"StackPushFloat": "StackPushFloat",
|
|
"StackPushObject": "StackPushObject",
|
|
"StackPopImage": "StackPopImage",
|
|
"StackPopString": "StackPopString",
|
|
"StackPopInt": "StackPopInt",
|
|
"StackPopFloat": "StackPopFloat",
|
|
"StackPopObject": "StackPopObject",
|
|
}
|