alphard/Endless-Nodes
1
0
Fork 0
mirror of https://github.com/tusharbhutt/Endless-Nodes.git synced 2026-03-21 20:42:12 -03:00
Code Issues Packages Projects Releases Wiki Activity
Files
fc9f488bc5f49cfbe1dbc121fef58a97db0788a6
Endless-Nodes/endless_nodes.py
tusharbhutt dcfbe32862 String/Numeric to numeric/string nodes fixed
2023-10-05 00:26:30 -06:00

1180 lines
39 KiB
Python
Raw Blame History

"""
@author: BiffMunky
@title: Endless ️🌊✨ Nodes
@nickname: ♾️🌊✨
@description: A small set of nodes I created for various numerical and text inputs. Features image saver with ability to have JSON saved to separate folder, parameter collection nodes, two aesthetic scoring models, switches for text and numbers, and conversion of string to numeric and vice versa.
"""
#----------------------------------------------
# Endless Sea of Stars Custom Node Collection
#https://github.com/tusharbhutt/Endless-Nodes
#----------------------------------------------
#
# Oct 04/23, V0.31: Release of V0.28 functionality (int, float, num to X), added String to X, code cleanup, vanity node renaming and recategorization
# Oct 04/23, V0.30: Squished bugs in the various X to X nodes
# Oct 03/23, V0.29: Save Image module added, saves images and JSON to separate folder if requested
# Sep 28/23, V0.28: (UNRELEASED) Added Variable types to X
# Sep 28/23, V0.27: (UNRELEASED) Corrected scoring nodes to actually add the value of the score into the image metadata .... still goobered!
# Sep 24/23, V0.26: (UNRELEASED) starting to correct scoring to get to image metadata
# Sep 24/23, V0.25: Added various X to String Nodes
# Sep 24/23, V0.24: Added In Image Reward scoring model with a single node to load model and output standard deviation and scoring via number or string nodes
# Sep 24/23, V0.23: Rework Aesthetic Score model and integrate it into single node to display score, added a requirements file
# Sep 23/23, V0.22: (UNRELEASED) Convert ImageReward output to base ten score
# Sep 22/23, V0.21: (UNRELEASED) Introduced aestheticscore, recategorized nodes into submenus, added some vanity coding to the node names, changed the ComfyUI manager header text
# Sep 21/23, V0.20: (UNRELEASED) Skeleton for save image
# Sep 21/23, V0.19: (UNRELEASED) Attempt for basic display nodes
# Sep 20/23, V0.16: Added Eight Input Number String
# Sep 18/23, V0.15: Added Combo Parameterizers to reduce number of nodes, allows for common resolution parameters to go to both pos/neg CLIP encode and adds separate pos/neg aesthetic score. Also has a version with pos/neg prompts
# Sep 18/23, V0.13: Fixed typos, added Paramaterizer with Prompt (unreleased to GitHub)
# Sep 18/23, V0.12: Added "Parameterizer", allows for parameters to be added to CLIP Encode
# Sep 15/23, V0.10: Added Six Input Number Widget, first release to GitHub
# Sep 12/23, V0.05: Added Six Input Number String
# Sep 08/23, V0.00: Basic Flow for Six Input Text Switch
#______________________________________________________________________________________________________________________________________________________________
# IMPORT MODULES BLOCK #
from PIL import Image
from PIL.PngImagePlugin import PngInfo
from colorama import init, Fore, Back, Style
from os.path import join
from warnings import filterwarnings
import ImageReward as RM
import clip
import colorama
import datetime
import folder_paths
import io
import json
import math
import numpy as np
import os
import pytorch_lightning as pl
import re
import socket
import statistics
import sys
import time
import torch
import torch.nn as nn
sys.path.insert(0, os.path.join(os.path.dirname(os.path.realpath(__file__)), "comfy"))
import comfy.sd
import comfy.utils
#import folder_paths
import typing as tg
# Initialize colorama for colored text
colorama.init(autoreset=True)
#______________________________________________________________________________________________________________________________________________________________
# "SWITCHES" BLOCK #
#
#----------------------------------------------
# Six Text Input Node for selection
class EndlessNode_SixTextInputSwitch:
def __init__(self):
pass
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"Input": ("INT", {"default": 1, "min": 1, "max": 6, "step": 1, "display": "slider"}),
#I like the slider idea, it's better for a touch screen
"text1": ("STRING", {"forceInput": True}),
},
"optional": {
"text2": ("STRING", {"forceInput": True}),
"text3": ("STRING", {"forceInput": True}),
"text4": ("STRING", {"forceInput": True}),
"text5": ("STRING", {"forceInput": True}),
"text6": ("STRING", {"forceInput": True}),
}
}
RETURN_TYPES = ("STRING",)
RETURN_NAMES = ("Output",)
FUNCTION = "six_text_switch"
CATEGORY = "Endless 🌊✨/Switches"
def six_text_switch(self, Input, text1=None,text2=None,text3=None,text4=None,text5=None,text6=None):
if Input == 1:
return (text1,)
elif Input == 2:
return (text2,)
elif Input == 3:
return (text3,)
elif Input == 4:
return (text4,)
elif Input == 5:
return (text5,)
else:
return (text6,)
#----------------------------------------------
# Eight Text Input Node for selection (needed more slots, what can I say)
class EndlessNode_EightTextInputSwitch:
def __init__(self):
pass
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"Input": ("INT", {"default": 1, "min": 1, "max": 8, "step": 1, "display": "slider"}),
#I like the slider idea, it's better for a touch screen
"text1": ("STRING", {"forceInput": True}),
},
"optional": {
"text2": ("STRING", {"forceInput": True}),
"text3": ("STRING", {"forceInput": True}),
"text4": ("STRING", {"forceInput": True}),
"text5": ("STRING", {"forceInput": True}),
"text6": ("STRING", {"forceInput": True}),
"text7": ("STRING", {"forceInput": True}),
"text8": ("STRING", {"forceInput": True}),
}
}
RETURN_TYPES = ("STRING",)
RETURN_NAMES = ("Output",)
FUNCTION = "eight_text_switch"
CATEGORY = "Endless 🌊✨/Switches"
def eight_text_switch(self,Input,text1=None,text2=None,text3=None,text4=None,text5=None,text6=None,text7=None,text8=None,):
if Input == 1:
return (text1,)
elif Input == 2:
return (text2,)
elif Input == 3:
return (text3,)
elif Input == 4:
return (text4,)
elif Input == 5:
return (text5,)
elif Input == 6:
return (text6,)
elif Input == 7:
return (text7,)
else:
return (text8,)
#----------------------------------------------
# Six Integer Input and Output via connectors
class EndlessNode_SixIntIOSwitch:
def __init__(self):
pass
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"INT1": ("INT", {"forceInput": True}),
},
"optional": {
"INT2": ("INT", {"forceInput": True}),
"INT3": ("INT", {"forceInput": True}),
"INT4": ("INT", {"forceInput": True}),
"INT5": ("INT", {"forceInput": True}),
"INT6": ("INT", {"forceInput": True}),
}
}
RETURN_TYPES = ("INT","INT","INT","INT","INT","INT",)
RETURN_NAMES = ("INT1","INT2","INT3","INT4","INT5","INT6",)
FUNCTION = "six_intIO_switch"
CATEGORY = "Endless 🌊✨/Switches"
def six_intIO_switch(self, Input, INT1=0, INT2=0, INT3=0, INT4=0, INT5=0, INT6=0):
if Input == 1:
return (INT1,)
elif Input == 2:
return (INT2,)
elif Input == 3:
return (INT3,)
elif Input == 4:
return (INT4,)
elif Input == 5:
return (INT5,)
else:
return (INT6,)
#----------------------------------------------
# Six Integer Input and Output by Widget
class EndlessNode_SixIntIOWidget:
def __init__(self):
pass
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"int1": ("INT", {"default": 0,}),
},
"optional": {
"int2": ("INT", {"default": 0,}),
"int3": ("INT", {"default": 0,}),
"int4": ("INT", {"default": 0,}),
"int5": ("INT", {"default": 0,}),
"int6": ("INT", {"default": 0,}),
}
}
RETURN_TYPES = ("INT","INT","INT","INT","INT","INT",)
RETURN_NAMES = ("INT1","INT2","INT3","INT4","INT5","INT6",)
FUNCTION = "six_int_widget"
CATEGORY = "Endless 🌊✨/Switches"
def six_int_widget(self,int1,int2,int3,int4,int5,int6):
return(int1,int2,int3,int4,int5,int6)
#______________________________________________________________________________________________________________________________________________________________
# PARAMETERS BLOCK #
#----------------------------------------------
# Text Encode Combo Box with prompt
class EndlessNode_XLParameterizerPrompt:
def __init__(self):
pass
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"base_width": ("INT", {"default": 1024, "min": 64, "max": 8192, "step": 16}),
"base_height": ("INT", {"default": 1024, "min": 64, "max": 8192, "step": 16}),
"base_crop_w": ("INT", {"default": 0, "min": 0, "max": 1024, "step": 8}),
"base_crop_h": ("INT", {"default": 0, "min": 0, "max": 1024, "step": 8}),
"base_target_w": ("INT", {"default": 1024, "min": 64, "max": 8192, "step": 16}),
"base_target_h": ("INT", {"default": 1024, "min": 64, "max": 8192, "step": 16}),
"refiner_width": ("INT", {"default": 1024, "min": 64, "max": 8192, "step": 16}),
"refiner_height": ("INT", {"default": 1024, "min": 64, "max": 8192, "step": 16}),
"refiner_ascore": ("FLOAT", {"default": 6, "min": 0, "max": 0xffffffffffffffff}),
},
"optional": {
"endlessG": ("STRING", {"default": "TEXT_G,acts as main prompt and connects to refiner text input", "multiline": True}),
"endlessL": ("STRING", {"default": "TEXT_L, acts as supporting prompt", "multiline": True}),
}
}
RETURN_TYPES = ("INT","INT","INT","INT","INT","INT","INT","INT","FLOAT","STRING","STRING",)
RETURN_NAMES = ("Base Width","Base Height","Base Cropped Width","Base Cropped Height","Base Target Width","Base Target Height","Refiner Width","Refiner Height","Refiner Aesthetic Score","Text_G/Refiner Prompt","Text_L Prompt",)
FUNCTION = "ParameterizerPrompt"
CATEGORY = "Endless 🌊✨/Parameters"
def ParameterizerPrompt(self,base_width,base_height,base_crop_w,base_crop_h,base_target_w,base_target_h,refiner_width,refiner_height,refiner_ascore,endlessG,endlessL):
return(base_width,base_height,base_crop_w,base_crop_h,base_target_w,base_target_h,refiner_width,refiner_height,refiner_ascore,endlessG,endlessL)
#----------------------------------------------
# CLIP text encode box without prompt
class EndlessNode_XLParameterizer:
def __init__(self):
pass
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"base_width": ("INT", {"default": 1024, "min": 64, "max": 8192, "step": 16}),
"base_height": ("INT", {"default": 1024, "min": 64, "max": 8192, "step": 16}),
"base_crop_w": ("INT", {"default": 0, "min": 0, "max": 8192, "step": 16}),
"base_crop_h": ("INT", {"default": 0, "min": 0, "max": 8192, "step": 16}),
"base_target_w": ("INT", {"default": 1024, "min": 64, "max": 8192, "step": 16}),
"base_target_h": ("INT", {"default": 1024, "min": 64, "max": 8192, "step": 16}),
"refiner_width": ("INT", {"default": 1024, "min": 64, "max": 8192, "step": 16}),
"refiner_height": ("INT", {"default": 1024, "min": 64, "max": 8192, "step": 16}),
"refiner_ascore": ("FLOAT", {"default": 6, "min": 0, "max": 0xffffffffffffffff}),
}
}
RETURN_TYPES = ("INT","INT","INT","INT","INT","INT","INT","INT","FLOAT",)
RETURN_NAMES = ("Base Width","Base Height","Base Cropped Width","Base Cropped Height","Base Target Width","Base Target Height","Refiner Width","Refiner Height","Refiner Aesthetic Score",)
FUNCTION = "Parameterizer"
CATEGORY = "Endless 🌊✨/Parameters"
def Parameterizer(self,base_width,base_height,base_crop_w,base_crop_h,base_target_w,base_target_h,refiner_width,refiner_height,refiner_ascore):
return(base_width,base_height,base_crop_w,base_crop_h,base_target_w,base_target_h,refiner_width,refiner_height,refiner_ascore)
#----------------------------------------------
# Text Encode Combo Box with prompt
class EndlessNode_ComboXLParameterizerPrompt:
def __init__(self):
pass
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"base_width": ("INT", {"default": 1024, "min": 64, "max": 8192, "step": 16}),
"base_height": ("INT", {"default": 1024, "min": 64, "max": 8192, "step": 16}),
"base_crop_w": ("INT", {"default": 0, "min": 0, "max": 8192, "step": 16}),
"base_crop_h": ("INT", {"default": 0, "min": 0, "max": 8192, "step": 16}),
"base_target_w": ("INT", {"default": 1024, "min": 64, "max": 8192, "step": 16}),
"base_target_h": ("INT", {"default": 1024, "min": 64, "max": 8192, "step": 16}),
"refiner_width": ("INT", {"default": 1024, "min": 64, "max": 8192, "step": 16}),
"refiner_height": ("INT", {"default": 1024, "min": 64, "max": 8192, "step": 16}),
"refiner_pascore": ("FLOAT", {"default": 6.5, "min": 0, "max": 0xffffffffffffffff}),
"refiner_nascore": ("FLOAT", {"default": 2.5, "min": 0, "max": 0xffffffffffffffff}),
},
"optional": {
"PendlessG": ("STRING", {"default": "Positive TEXT_G,acts as main prompt and connects to refiner text input", "multiline": True}),
"PendlessL": ("STRING", {"default": "Positive TEXT_L, acts as supporting prompt", "multiline": True}),
"NendlessG": ("STRING", {"default": "Negative TEXT_G, acts as main prompt and connects to refiner text input", "multiline": True}),
"NendlessL": ("STRING", {"default": "Negative TEXT_L, acts as supporting prompt", "multiline": True}),
}
}
RETURN_TYPES = ("INT","INT","INT","INT","INT","INT","INT","INT","FLOAT","FLOAT","STRING","STRING", "STRING","STRING",)
RETURN_NAMES = ("Base Width","Base Height","Base Cropped Width","Base Cropped Height","Base Target Width","Base Target Height","Refiner Width","Refiner Height","Positive Refiner Aesthetic Score","Negative Refiner Aesthetic Score","Positive Text_G and Refiner Text Prompt","Postive Text_L Prompt","Negative Text_G and Refiner Text Prompt","Negative Text_L Prompt",)
FUNCTION = "ComboParameterizerPrompt"
CATEGORY = "Endless 🌊✨/Parameters"
def ComboParameterizerPrompt(self,base_width,base_height,base_crop_w,base_crop_h,base_target_w,base_target_h,refiner_width,refiner_height,refiner_pascore,refiner_nascore,PendlessG,PendlessL,NendlessG,NendlessL):
return(base_width,base_height,base_crop_w,base_crop_h,base_target_w,base_target_h,refiner_width,refiner_height,refiner_pascore,refiner_nascore,PendlessG,PendlessL,NendlessG,NendlessL)
#----------------------------------------------
# CLIP text encode box without prompt, COMBO that allows one box for both pos/neg parameters to be fed to CLIP text, with separate POS/NEG Aesthetic score
class EndlessNode_ComboXLParameterizer:
def __init__(self):
pass
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"base_width": ("INT", {"default": 1024, "min": 64, "max": 8192, "step": 16}),
"base_height": ("INT", {"default": 1024, "min": 64, "max": 8192, "step": 16}),
"base_crop_w": ("INT", {"default": 0, "min": 0, "max": 8192, "step": 16}),
"base_crop_h": ("INT", {"default": 0, "min": 0, "max": 8192, "step": 16}),
"base_target_w": ("INT", {"default": 1024, "min": 64, "max": 8192, "step": 16}),
"base_target_h": ("INT", {"default": 1024, "min": 64, "max": 8192, "step": 16}),
"refiner_width": ("INT", {"default": 1024, "min": 64, "max": 8192, "step": 16}),
"refiner_height": ("INT", {"default": 1024, "min": 64, "max": 8192, "step": 16}),
"refiner_pascore": ("FLOAT", {"default": 6.5, "min": 0, "max": 0xffffffffffffffff}),
"refiner_nascore": ("FLOAT", {"default": 2.5, "min": 0, "max": 0xffffffffffffffff}),
}
}
RETURN_TYPES = ("INT","INT","INT","INT","INT","INT","INT","INT","FLOAT","FLOAT")
RETURN_NAMES = ("Base Width","Base Height","Base Cropped Width","Base Cropped Height","Base Target Width","Base Target Height","Refiner Width","Refiner Height","Positive Refiner Aesthetic Score","Negative Refiner Aesthetic Score",)
FUNCTION = "ComboParameterizer"
CATEGORY = "Endless 🌊✨/Parameters"
def ComboParameterizer(self,base_width,base_height,base_crop_w,base_crop_h,base_target_w,base_target_h,refiner_width,refiner_height,refiner_pascore, refiner_nascore):
return(base_width,base_height,base_crop_w,base_crop_h,base_target_w,base_target_h,refiner_width,refiner_height,refiner_pascore, refiner_nascore)
#______________________________________________________________________________________________________________________________________________________________
# IMAGE SCORING BLOCK #
#----------------------------------------------
# Aesthetic Scoring Node
folder_paths.folder_names_and_paths["aesthetic"] = ([os.path.join(folder_paths.models_dir,"aesthetic")], folder_paths.supported_pt_extensions)
class MLP(pl.LightningModule):
def __init__(self, input_size, xcol='emb', ycol='avg_rating'):
super().__init__()
self.input_size = input_size
self.xcol = xcol
self.ycol = ycol
self.layers = nn.Sequential(
nn.Linear(self.input_size, 1024),
#nn.ReLU(),
nn.Dropout(0.2),
nn.Linear(1024, 128),
#nn.ReLU(),
nn.Dropout(0.2),
nn.Linear(128, 64),
#nn.ReLU(),
nn.Dropout(0.1),
nn.Linear(64, 16),
#nn.ReLU(),
nn.Linear(16, 1)
)
def forward(self, x):
return self.layers(x)
def training_step(self, batch, batch_idx):
x = batch[self.xcol]
y = batch[self.ycol].reshape(-1, 1)
x_hat = self.layers(x)
loss = F.mse_loss(x_hat, y)
return loss
def validation_step(self, batch, batch_idx):
x = batch[self.xcol]
y = batch[self.ycol].reshape(-1, 1)
x_hat = self.layers(x)
loss = F.mse_loss(x_hat, y)
return loss
def configure_optimizers(self):
optimizer = torch.optim.Adam(self.parameters(), lr=1e-3)
return optimizer
def normalized(a, axis=-1, order=2):
import numpy as np # pylint: disable=import-outside-toplevel
l2 = np.atleast_1d(np.linalg.norm(a, order, axis))
l2[l2 == 0] = 1
return a / np.expand_dims(l2, axis)
class EndlessNode_Scoring:
def __init__(self):
pass
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"model_name": (folder_paths.get_filename_list("aesthetic"), {"multiline": False, "default": "chadscorer.pth"}),
"image": ("IMAGE",),
}
}
RETURN_TYPES = ("NUMBER","IMAGE")
FUNCTION = "calc_score"
CATEGORY = "Endless 🌊✨/Scoring"
def calc_score(self, model_name, image):
m_path = folder_paths.folder_names_and_paths["aesthetic"][0]
m_path2 = os.path.join(m_path[0], model_name)
model = MLP(768) # CLIP embedding dim is 768 for CLIP ViT L 14
s = torch.load(m_path2)
model.load_state_dict(s)
model.to("cuda")
model.eval()
device = "cuda"
model2, preprocess = clip.load("ViT-L/14", device=device) # RN50x64
tensor_image = image[0]
img = (tensor_image * 255).to(torch.uint8).numpy()
pil_image = Image.fromarray(img, mode='RGB')
image2 = preprocess(pil_image).unsqueeze(0).to(device)
with torch.no_grad():
image_features = model2.encode_image(image2)
im_emb_arr = normalized(image_features.cpu().detach().numpy())
prediction = model(torch.from_numpy(im_emb_arr).to(device).type(torch.cuda.FloatTensor))
final_prediction = round(float(prediction[0]), 2)
del model
return (final_prediction,)
# #---------------------------------------------- NOT WORKING, NEED TO LOOK AT IT
# # Aesthetic Scoring Node with Scoring passed to image
# class EndlessNode_ScoringAutoScore:
# def __init__(self):
# pass
# @classmethod
# def INPUT_TYPES(cls):
# return {
# "required": {
# "model_name": (folder_paths.get_filename_list("aesthetic"), {"multiline": False, "default": "chadscorer.pth"}),
# "image": ("IMAGE",),
# }
# }
# RETURN_TYPES = ("NUMBER","IMAGE")
# FUNCTION = "calc_score"
# OUTPUT_NODE = True
# CATEGORY = "Endless 🌊✨/Scoring"
# def calc_score(self, model_name, image):
# m_path = folder_paths.folder_names_and_paths["aesthetic"][0]
# m_path2 = os.path.join(m_path[0], model_name)
# model = MLP(768) # CLIP embedding dim is 768 for CLIP ViT L 14
# s = torch.load(m_path2)
# model.load_state_dict(s)
# model.to("cuda")
# model.eval()
# device = "cuda"
# model2, preprocess = clip.load("ViT-L/14", device=device) # RN50x64
# tensor_image = image[0]
# img = (tensor_image * 255).to(torch.uint8).numpy()
# pil_image = Image.fromarray(img, mode='RGB')
# image2 = preprocess(pil_image).unsqueeze(0).to(device)
# with torch.no_grad():
# image_features = model2.encode_image(image2)
# im_emb_arr = normalized(image_features.cpu().detach().numpy())
# prediction = model(torch.from_numpy(im_emb_arr).to(device).type(torch.cuda.FloatTensor))
# final_prediction = round(float(prediction[0]), 2)
# del model
# # Metadata part
# extra_pnginfo = {"SCORE": str(final_prediction)}
# return (final_prediction, image)
#----------------------------------------------
# Image Reward Scoring
class EndlessNode_ImageReward:
def __init__(self):
self.model = None
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"model": ("STRING", {"multiline": False, "default": "ImageReward-v1.0"}),
"prompt": ("STRING", {"multiline": True, "forceInput": True}),
"images": ("IMAGE",),
},
}
RETURN_TYPES = ("FLOAT", "STRING", "FLOAT", "STRING")
RETURN_NAMES = ("SCORE_FLOAT", "SCORE_STRING", "VALUE_FLOAT", "VALUE_STRING")
CATEGORY = "Endless 🌊✨/Scoring"
FUNCTION = "process_images"
def process_images(self, model, prompt, images,): #rounded):
if self.model is None:
self.model = RM.load(model)
score = 0.0
for image in images:
# convert to PIL image
i = 255.0 * image.cpu().numpy()
img = Image.fromarray(np.clip(i, 0, 255).astype(np.uint8))
score += self.model.score(prompt, [img])
score /= len(images)
# assume std dev follows normal distribution curve
valuescale = 0.5 * (1 + math.erf(score / math.sqrt(2))) * 10 # *10 to get a value between -10
return (score, str(score), valuescale, str(valuescale))
#---------------------------------------------- NOT WORKING, NEED TO LOOK AT
# Image Reward Scoring with score passed to image
class EndlessNode_ImageRewardAutoScore:
def __init__(self):
self.model = None
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"model": ("STRING", {"multiline": False, "default": "ImageReward-v1.0"}),
"prompt": ("STRING", {"multiline": True, "forceInput": True}),
"images": ("IMAGE",),
},
}
RETURN_TYPES = ("FLOAT", "STRING", "FLOAT", "STRING", "IMAGE")
RETURN_NAMES = ("SCORE_FLOAT", "SCORE_STRING", "VALUE_FLOAT", "VALUE_STRING", "TO_IMAGE")
OUTPUT_NODE = True
CATEGORY = "Endless 🌊✨/Scoring"
FUNCTION = "score_meta"
def score_meta(self, model, prompt, images):
if self.model is None:
self.model = RM.load(model)
# Scoring part
score = 0.0
for image in images:
i = 255.0 * image.cpu().numpy()
img = Image.fromarray(np.clip(i, 0, 255).astype(np.uint8))
score += self.model.score(prompt, [img])
score /= len(images)
valuescale = 0.5 * (1 + math.erf(score / math.sqrt(2))) * 10
# Metadata part
extra_pnginfo = {"SCORE": str(score)}
# Returning both the score and the modified image
return (score, str(score), valuescale, str(valuescale), images)
# ______________________________________________________________________________________________________________________________________________________________
# IMAGE SAVERS BLOCK #
# ----------------------------------------------
# Saver type one: saves IMAGE and JSON files, can specify separate folders for each, or one, or none, and use Python timestamps
class EndlessNode_ImageSaver:
def __init__(self):
self.output_dir = folder_paths.get_output_directory()
self.type = "output"
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"images": ("IMAGE",),
"filename_prefix": ("STRING", {"default": "ComfyUI"}),
"delimiter": ("STRING", {"default": "_"}),
"filename_number_padding": ("INT", {"default": 4, "min": 1, "max": 9, "step": 1}),
"filename_number_start": (["false", "true"],),
"img_folder": ("STRING", {"default": None}),
"json_folder": ("STRING", {"default": None}),
},
"hidden": {
"prompt": "PROMPT", "extra_pnginfo": "EXTRA_PNGINFO"
},
}
RETURN_TYPES = ()
FUNCTION = "save_images"
OUTPUT_NODE = True
CATEGORY = "Endless 🌊✨/IO"
def save_images(self, images, filename_prefix="ComfyUI", delimiter="_",
filename_number_padding=4, filename_number_start='false',
img_folder=None, json_folder=None, prompt=None, extra_pnginfo=None):
# Replace illegal characters in the filename prefix with dashes
filename_prefix = re.sub(r'[<>:"\/\\|?*]', '-', filename_prefix)
# Set IMG Extension
img_extension = '.png'
counter = 1
results = list()
for image in images:
i = 255. * image.cpu().numpy()
img = Image.fromarray(np.clip(i, 0, 255).astype(np.uint8))
metadata = PngInfo()
if prompt is not None:
metadata.add_text("prompt", json.dumps(prompt))
if extra_pnginfo is not None:
for x in extra_pnginfo:
metadata.add_text(x, json.dumps(extra_pnginfo[x]))
img_file, json_file = self.generate_filenames(filename_prefix, delimiter, counter,filename_number_padding, filename_number_start,img_extension, img_folder, json_folder)
try:
if img_extension == '.png':
img.save(img_file, pnginfo=metadata, compress_level=4)
elif img_extension == '.jpeg':
img.save(img_file, quality=100, optimize=True)
with open(json_file, 'w', encoding='utf-8', newline='\n') as f:
if prompt is not None:
f.write(json.dumps(prompt, indent=4))
print(Fore.GREEN + f"+ File(s) saved to: {img_file}")
results.append({
"image_filename": os.path.basename(img_file),
"image_path": img_file,
"json_filename": os.path.basename(json_file),
"text_path": json_file,
"type": self.type
})
except OSError as e:
print(Fore.RED + " + Unable to save file: ", end='')
print({img_file})
print(e)
except Exception as e:
print(Fore.RED + " + Unable to save file due to the following error: ", end='')
print(e)
counter += 1
return {"ui": {"results": results}}
def generate_filenames(self, filename_prefix, delimiter, counter, filename_number_padding, filename_number_start, img_extension, img_folder, json_folder):
if filename_number_start == 'true':
img_file = f"{filename_prefix}{delimiter}{counter:0{filename_number_padding}}{img_extension}"
json_file = f"{filename_prefix}{delimiter}{counter:0{filename_number_padding}}.json"
else:
img_file = f"{filename_prefix}{delimiter}{counter:0{filename_number_padding}}{img_extension}"
json_file = f"{filename_prefix}{delimiter}{counter:0{filename_number_padding}}.json"
# Construct full paths for image and text files based on folders provided
if img_folder:
img_folder = self.replace_date_time_placeholders(img_folder)
os.makedirs(img_folder, exist_ok=True) # Create the image folder if it doesn't exist
img_file = os.path.join(img_folder, img_file)
else:
img_file = os.path.join(self.output_dir, img_file)
if json_folder:
json_folder = self.replace_date_time_placeholders(json_folder)
os.makedirs(json_folder, exist_ok=True) # Create the image folder if it doesn't exist
json_file = os.path.join(json_folder, json_file)
else:
json_file = os.path.join(os.path.dirname(img_file), json_file)
# Apply placeholders for date and time in filenames and folder
img_file = self.replace_date_time_placeholders(img_file)
json_file = self.replace_date_time_placeholders(json_file)
# Check if files with the same name exist, increment counter if necessary
while os.path.exists(img_file) or os.path.exists(json_file):
counter += 1
if filename_number_start == 'true':
img_file = f"{filename_prefix}{delimiter}{counter:0{filename_number_padding}}{img_extension}"
json_file = f"{filename_prefix}{delimiter}{counter:0{filename_number_padding}}.json"
else:
img_file = f"{filename_prefix}{delimiter}{counter:0{filename_number_padding}}{img_extension}"
json_file = f"{filename_prefix}{delimiter}{counter:0{filename_number_padding}}.json"
# Construct full paths for image and text files based on folders provided
if img_folder:
img_folder = self.replace_date_time_placeholders(img_folder)
os.makedirs(img_folder, exist_ok=True) # Create the image folder if it doesn't exist
img_file = os.path.join(img_folder, img_file)
else:
img_file = os.path.join(self.output_dir, img_file)
if json_folder:
json_folder = self.replace_date_time_placeholders(json_folder)
os.makedirs(json_folder, exist_ok=True) # Create the image folder if it doesn't exist
json_file = os.path.join(json_folder, json_file)
else:
json_file = os.path.join(os.path.dirname(img_file), json_file)
# Apply placeholders for date and time in filenames and folder
img_file = self.replace_date_time_placeholders(img_file)
json_file = self.replace_date_time_placeholders(json_file)
return img_file, json_file
def replace_date_time_placeholders(self, filename):
# Replace date and time placeholders with actual date and time strings
now = datetime.datetime.now()
placeholders = {
'%Y': now.strftime('%Y'), # Year with century as a decimal number
'%y': now.strftime('%y'), # Year without century as a zero-padded decimal number
'%m': now.strftime('%m'), # Month as a zero-padded decimal number
'%d': now.strftime('%d'), # Day of the month as a zero-padded decimal number
'%H': now.strftime('%H'), # Hour (24-hour clock) as a zero-padded decimal number
'%M': now.strftime('%M'), # Minute as a zero-padded decimal number
'%S': now.strftime('%S'), # Second as a zero-padded decimal number
}
for placeholder, replacement in placeholders.items():
filename = filename.replace(placeholder, replacement)
return filename
# ______________________________________________________________________________________________________________________________________________________________
# CONVERTER NODES BLOCK #
#
# ----------------------------------------------
# Float to Integer
class EndlessNode_FloattoInt:
CATEGORY = "Endless 🌊✨/Converters/Float"
def __init__(self):
pass
@classmethod
def INPUT_TYPES(cls):
return {
"required": {"FloatValue": ("FLOAT", {"default": 0.0})},
}
RETURN_TYPES = ("INT",)
FUNCTION = "inputfloat"
def inputfloat(self, FloatValue):
return (int(FloatValue),)
# ----------------------------------------------
# Float value to Number. There is no real "Number" function in Python, this is here so that nodes that need a NUMBER can take the FLOAT value
class EndlessNode_FloattoNum:
CATEGORY = "Endless 🌊✨/Converters/Float"
def __init__(self):
pass
@classmethod
def INPUT_TYPES(cls):
return {
"required": {"FloatValue": ("FLOAT", {"default": 0.0})},
}
RETURN_TYPES = ("NUMBER",)
FUNCTION = "inputfloat"
def inputfloat(self, FloatValue):
return (float(FloatValue),)
# ----------------------------------------------
# Float to String,
class EndlessNode_FloattoString:
CATEGORY = "Endless 🌊✨/Converters/Float"
def __init__(self):
pass
@classmethod
def INPUT_TYPES(cls):
return {
"required": {"FloatValue": ("FLOAT", {"default": 0.0})},
}
RETURN_TYPES = ("STRING")
FUNCTION = "inputfloat"
def inputfloat(self, FloatValue):
return(str(FloatValue),)
# ----------------------------------------------
# Float value to X
class EndlessNode_FloattoX:
CATEGORY = "Endless 🌊✨/Converters/Float"
def __init__(self):
pass
@classmethod
def INPUT_TYPES(cls):
return {
"required": {"FloatValue": ("FLOAT", {"default": 0.0})},
}
RETURN_TYPES = ("INT", "NUMBER","STRING")
FUNCTION = "inputfloat"
def inputfloat(self, FloatValue):
return(int(FloatValue), float(FloatValue),str(FloatValue),)
# ----------------------------------------------
# Integer to Float
class EndlessNode_InttoFloat:
CATEGORY = "Endless 🌊✨/Converters/Int"
def __init__(self):
pass
@classmethod
def INPUT_TYPES(cls):
return {
"required": {"IntegerValue": ("INT", {"default": 0})},
}
RETURN_TYPES = ("FLOAT",)
FUNCTION = "inputint"
def inputint(self, IntegerValue):
return (int(IntegerValue),)
# ----------------------------------------------
# Integer to Number (for compatability purposes)
class EndlessNode_InttoNum:
CATEGORY = "Endless 🌊✨/Converters/Int"
def __init__(self):
pass
@classmethod
def INPUT_TYPES(cls):
return {
"required": {"IntegerValue": ("INT", {"default": 0})},
}
RETURN_TYPES = ("NUMBER")
FUNCTION = "inputint"
def inputint(self, IntegerValue):
return (int(IntegerValue),)
# ----------------------------------------------
# Integer to String
class EndlessNode_InttoString:
CATEGORY = "Endless 🌊✨/Converters/Int"
@classmethod
def INPUT_TYPES(cls):
return {
"required": {"IntegerValue": ("INT", {"default": 0})},
}
RETURN_TYPES = ("STRING")
FUNCTION = "inputint"
def inputint(self, IntegerValue):
return (str(IntegerValue),)
# ----------------------------------------------
#Integer to X
class EndlessNode_InttoX:
CATEGORY = "Endless 🌊✨/Converters/Int"
def __init__(self):
pass
@classmethod
def INPUT_TYPES(cls):
return {
"required": {"IntegerValue": ("INT", {"default": 0})},
}
RETURN_TYPES = ("FLOAT", "NUMBER","STRING")
FUNCTION = "inputint"
def inputint(self, IntegerValue):
return(float(IntegerValue), float(IntegerValue),str(IntegerValue),)
# ----------------------------------------------
# Number to Float
class EndlessNode_NumtoFloat:
CATEGORY = "Endless 🌊✨/Converters/Number"
def __init__(self):
pass
@classmethod
def INPUT_TYPES(cls):
return {
"required": {"NumberValue": ("NUMBER",),}
}
RETURN_TYPES = ("FLOAT",)
FUNCTION = "inputnum"
def inputnum(self, NumberValue):
return (float(NumberValue),)
# ----------------------------------------------
# Number to Integer
class EndlessNode_NumtoInt:
CATEGORY = "Endless 🌊✨/Converters/Number"
def __init__(self):
pass
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"number": ("NUMBER",),}
}
RETURN_TYPES = ("INT",)
FUNCTION = "number_to_int"
def number_to_int(self, number):
return (int(number), )
# ----------------------------------------------
# Number to String
class EndlessNode_NumtoString:
def __init__(self):
pass
CATEGORY = "Endless 🌊✨/Converters/Number"
@classmethod
@classmethod
def INPUT_TYPES(cls):
return {
"required": {"NumberValue": ("NUMBER",)},
}
RETURN_TYPES = ("STRING")
FUNCTION = "inputnum"
def inputnum(self, NumberValue):
return(str(NumberValue),)
# ----------------------------------------------
# Number to X
class EndlessNode_NumtoX:
CATEGORY = "Endless 🌊✨/Converters/Number"
def __init__(self):
pass
@classmethod
def INPUT_TYPES(cls):
return {
"required": {"NumberValue": ("NUMBER",)},
}
RETURN_TYPES = ("FLOAT", "INT", "STRING")
RETURN_NAMES = ("FLOAT", "INT", "STRING")
FUNCTION = "inputnum"
def inputnum(self, NumberValue):
return(float(NumberValue), int(NumberValue),str(NumberValue),)
# ----------------------------------------------
# String to Float
class EndlessNode_StringtoFloat:
CATEGORY = "Endless 🌊✨/Converters/String"
def __init__(self):
pass
@classmethod
def INPUT_TYPES(cls):
return {
"required": {"StringValue": ("STRING", {"default": ""})},
}
RETURN_TYPES = ("FLOAT",)
FUNCTION = "inputstring"
def inputstring(self, StringValue):
try:
return(float(StringValue),)
except (ValueError, TypeError): # Handle non-numerical input here by returning default value of 0
return 0.0
# ----------------------------------------------
# String to Int
class EndlessNode_StringtoInt:
CATEGORY = "Endless 🌊✨/Converters/String"
def __init__(self):
pass
@classmethod
def INPUT_TYPES(cls):
return {
"required": {"StringValue": ("STRING", {"default": ""})},
}
RETURN_TYPES = ("INT",)
FUNCTION = "inputstring"
def inputstring(self, StringValue):
try:
return(int(float(StringValue)),) # int(float(x)) in case the user puts in a decimal
except (ValueError, TypeError): # Handle non-numerical input here by returning default value of 0
return 0
# ----------------------------------------------
# String to Number. There is no real "Number" function in Python, this is here so that nodes that need a NUMBER can take the FLOAT value
class EndlessNode_StringtoNum:
CATEGORY = "Endless 🌊✨/Converters/String"
def __init__(self):
pass
@classmethod
def INPUT_TYPES(cls):
return {
"required": {"StringValue": ("STRING", {"default": ""})},
}
RETURN_TYPES = ("NUMBER",)
FUNCTION = "inputstring"
def inputstring(self, StringValue):
try:
return(float(StringValue),)
except (ValueError, TypeError): # Handle non-numerical input here by returning default value of 0
return 0.0
# ----------------------------------------------
# String to X
class EndlessNode_StringtoX:
CATEGORY = "Endless 🌊✨/Converters/String"
def __init__(self):
pass
@classmethod
def INPUT_TYPES(cls):
return {
"required": {"StringValue": ("STRING", {"default": ""})},
}
RETURN_TYPES = ("INT", "FLOAT","NUMBER")
FUNCTION = "inputstring"
def inputstring(self, StringValue):
try:
return(int(float(StringValue)), float(StringValue),float(StringValue),) # int(float(x)) in case the user puts in a decimal
except (ValueError, TypeError): # Handle non-numerical input here by returning default value of 0
return 0, 0.0, 0.0
#______________________________________________________________________________________________________________________________________________________________
# CREDITS #
#
#
# Comfyroll Custom Nodes for the initial node code layout, coding snippets, and inspiration for the text input and number switches
#
# https://github.com/RockOfFire/ComfyUI_Comfyroll_CustomNode
#
# WLSH Nodes for some coding for the Integer Widget
# https://github.com/wallish77/wlsh_nodes
#
# ComfyUI Interface for the basic ideas of what nodes I wanted
#
# https://github.com/comfyanonymous/ComfyUI
#
# ComfyUI-Strimmlarns-Aesthetic-Score for the original coding for Aesthetic Scoring
#
# https://github.com/strimmlarn/ComfyUI-Strimmlarns-Aesthetic-Score
#
# The scorer uses the MLP class code from Christoph Schuhmann
#
#https://github.com/christophschuhmann/improved-aesthetic-predictor
#[Zane A's ComfyUI-ImageReward](https://github.com/ZaneA/ComfyUI-ImageReward) for the original coding for the Image Reward node
#
#Zane's node in turn uses [ImageReward](https://github.com/THUDM/ImageReward)
#
#
#Mikey nodes to grab code snippet to pass scoring metadata to image
#
#https://github.com/bash-j/mikey_nodes
#
# Took some base code from the WAS save image node to repurpose it
#
#https://github.com/WASasquatch/was-node-suite-comfyui
#--------------------------------------
#
# Special thanks to [chrisgoringe](https://github.com/chrisgoringe) for some vital insight into correcting messy commas in the tuples for the converter nodes, much appreciated!
#
#######################################################################################
# CELLAR DWELLERS
Reference in New Issue View Git Blame Copy Permalink