1064 lines
58 KiB
Python
1064 lines
58 KiB
Python
from ..components.tree import TREE_OUTPUTS
|
|
import os
|
|
import folder_paths
|
|
import re
|
|
import json
|
|
import time
|
|
import numpy as np
|
|
import pyexiv2
|
|
from PIL.PngImagePlugin import PngInfo
|
|
from PIL import Image
|
|
from pathlib import Path
|
|
import datetime
|
|
import comfy.samplers
|
|
from comfy import model_management
|
|
import random
|
|
import nodes
|
|
import comfy_extras.nodes_flux as nodes_flux
|
|
import torch
|
|
from ..components import utility
|
|
from ..components import primeresamplers
|
|
from ..components import file_output
|
|
from server import PromptServer
|
|
from ..components.tree import PRIMERE_ROOT
|
|
from comfy.cli_args import args
|
|
from .modules import exif_data_checker
|
|
from ..Nodes.Visuals import PrimereVisualCKPT
|
|
from ..Nodes.Visuals import PrimereVisualStyle
|
|
from transformers import pipeline
|
|
from torchvision.transforms import functional as TF
|
|
import comfy_extras.nodes_model_advanced as nodes_model_advanced
|
|
|
|
ALLOWED_EXT = file_output.ALLOWED_EXT
|
|
|
|
class PrimereMetaSave:
|
|
RETURN_TYPES = ("STRING",)
|
|
RETURN_NAMES = ("SAVED_INFO",)
|
|
FUNCTION = "save_images_meta"
|
|
OUTPUT_NODE = True
|
|
CATEGORY = TREE_OUTPUTS
|
|
|
|
NODE_FILE = os.path.abspath(__file__)
|
|
NODE_ROOT = os.path.dirname(NODE_FILE)
|
|
|
|
def __init__(self):
|
|
self.output_dir = folder_paths.output_directory
|
|
self.type = 'output'
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"save_image": ("BOOLEAN", {"default": True}),
|
|
"aesthetic_trigger": ("INT", {"default": 0, "min": 0, "max": 1000, "step": 1}),
|
|
"images": ("IMAGE",),
|
|
"output_path": ("STRING", {"default": '[time(%Y-%m-%d)]', "multiline": False}),
|
|
"subpath": (["None", "Dev", "Test", "Serie", "Production", "Preview", "NewModel", "Project", "Portfolio", "Civitai", "Behance", "Facebook", "Instagram", "Character", "Style", "Product", "Fun", "SFW", "NSFW"], {"default": "Project"}),
|
|
"subpath_priority": ("BOOLEAN", {"default": False, "label_on": "Preferred", "label_off": "Selected subpath"}),
|
|
"add_modelname_to_path": ("BOOLEAN", {"default": False}),
|
|
"add_concept_to_path": ("BOOLEAN", {"default": False}),
|
|
"filename_prefix": ("STRING", {"default": "ComfyUI"}),
|
|
"filename_delimiter": ("STRING", {"default": "_"}),
|
|
"add_date_to_filename": ("BOOLEAN", {"default": True}),
|
|
"add_time_to_filename": ("BOOLEAN", {"default": True}),
|
|
"add_seed_to_filename": ("BOOLEAN", {"default": True}),
|
|
"add_size_to_filename": ("BOOLEAN", {"default": True}),
|
|
"add_ascore_to_filename": ("BOOLEAN", {"default": True}),
|
|
"filename_number_padding": ("INT", {"default": 2, "min": 1, "max": 9, "step": 1}),
|
|
"filename_number_start": ("BOOLEAN", {"default":False}),
|
|
"extension": (['png', 'jpeg', 'jpg', 'gif', 'tiff', 'webp'], {"default": "jpg"}),
|
|
"png_embed_workflow": ("BOOLEAN", {"default": False}),
|
|
"png_embed_data": ("BOOLEAN", {"default": False}),
|
|
"image_embed_exif": ("BOOLEAN", {"default": False}),
|
|
"a1111_civitai_meta": ("BOOLEAN", {"default": False}),
|
|
"quality": ("INT", {"default": 95, "min": 1, "max": 100, "step": 1}),
|
|
"overwrite_mode": (["false", "prefix_as_filename"],),
|
|
"save_meta_to_json": ("BOOLEAN", {"default": False}),
|
|
"save_info_to_txt": ("BOOLEAN", {"default": False}),
|
|
},
|
|
"optional": {
|
|
"image_metadata": ('TUPLE', {"forceInput": True}),
|
|
},
|
|
"hidden": {
|
|
"prompt": "PROMPT",
|
|
"extra_pnginfo": "EXTRA_PNGINFO"
|
|
},
|
|
}
|
|
|
|
def save_images_meta(self, images, add_date_to_filename, add_time_to_filename, add_seed_to_filename, add_size_to_filename, add_ascore_to_filename, save_meta_to_json, save_info_to_txt, image_metadata=None,
|
|
output_path='[time(%Y-%m-%d)]', subpath='Project', subpath_priority = False, add_modelname_to_path = False, add_concept_to_path = False, filename_prefix="ComfyUI", filename_delimiter='_',
|
|
extension='jpg', quality=95, prompt=None, extra_pnginfo=None,
|
|
overwrite_mode='false', filename_number_padding=2, filename_number_start=False,
|
|
png_embed_workflow=False, png_embed_data=False, image_embed_exif=False, a1111_civitai_meta=False, save_image=True, aesthetic_trigger = 0):
|
|
|
|
if save_image == False:
|
|
saved_info = "*** Image saver switched OFF, image not saved. ***"
|
|
return saved_info, {"ui": {"images": []}}
|
|
|
|
if 'aesthetic_score' in image_metadata:
|
|
if (type(image_metadata['aesthetic_score']).__name__ == 'int'):
|
|
image_metadata['aesthetic_score'] = str(image_metadata['aesthetic_score'])
|
|
|
|
if (image_metadata['aesthetic_score'].isdigit()) and int(image_metadata['aesthetic_score']) > 0:
|
|
if aesthetic_trigger > int(image_metadata['aesthetic_score']):
|
|
saved_info = "*** Image ignored because aesthetic score: [" + str(image_metadata['aesthetic_score']) + "] less than trigger setting: [" + str(aesthetic_trigger) + "]. ***"
|
|
return saved_info, {"ui": {"images": []}}
|
|
|
|
delimiter = filename_delimiter
|
|
number_padding = filename_number_padding
|
|
tokens = file_output.TextTokens()
|
|
|
|
original_output = self.output_dir
|
|
filename_prefix = file_output.sanitize_path_part(tokens.parseTokens(filename_prefix))
|
|
nowdate = datetime.datetime.now()
|
|
if image_metadata is None:
|
|
image_metadata = {}
|
|
|
|
if len(images) < 1:
|
|
return
|
|
|
|
image_metadata['saved_image_width'] = images[0].shape[1]
|
|
image_metadata['saved_image_heigth'] = images[0].shape[0]
|
|
# image_metadata['upscaler_ratio'] = round(image_metadata['saved_image_width'] / image_metadata['width'], 2)
|
|
if 'width' in image_metadata and 'height' in image_metadata:
|
|
image_metadata['upscaler_ratio'] = 'From: ' + str(image_metadata['width']) + 'x' + str(image_metadata['height']) + ' to: ' + str(image_metadata['saved_image_width']) + 'x' + str(image_metadata['saved_image_heigth']) + ' Ratio: ' + str(round(round(image_metadata['saved_image_width'] / image_metadata['width'] / 0.05) * 0.05, 2))
|
|
|
|
if add_ascore_to_filename:
|
|
if 'aesthetic_score' in image_metadata:
|
|
if (image_metadata['aesthetic_score'].isdigit()) and int(image_metadata['aesthetic_score']) > 0:
|
|
filename_prefix = filename_prefix + '_A' + str(image_metadata['aesthetic_score'])
|
|
if add_date_to_filename:
|
|
filename_prefix = filename_prefix + '_' + nowdate.strftime("%Y%d%m")
|
|
if add_time_to_filename:
|
|
filename_prefix = filename_prefix + '_' + nowdate.strftime("%H%M%S")
|
|
if add_seed_to_filename:
|
|
if 'seed' in image_metadata:
|
|
filename_prefix = filename_prefix + '_' + str(image_metadata['seed'])
|
|
if add_size_to_filename:
|
|
if 'width' in image_metadata:
|
|
filename_prefix = filename_prefix + '_' + str(image_metadata['saved_image_width']) + 'x' + str(image_metadata['saved_image_heigth'])
|
|
|
|
output_path = file_output.parse_output_path_base(output_path, self.output_dir)
|
|
base_output = os.path.basename(output_path)
|
|
if output_path.endswith("ComfyUI/output") or output_path.endswith("ComfyUI\output"):
|
|
base_output = ""
|
|
|
|
subdirs = []
|
|
|
|
if add_concept_to_path == True and 'model_concept' in image_metadata:
|
|
concept_name = image_metadata['model_concept']
|
|
if concept_name == 'Auto' and 'model_version' in image_metadata:
|
|
match image_metadata['model_version']:
|
|
case 'SDXL_2048':
|
|
concept_name = 'SDXL'
|
|
case 'BaseModel_768':
|
|
concept_name = 'SD1'
|
|
case 'SD3_1024':
|
|
concept_name = 'SD3'
|
|
case 'Stable_Zero123_768':
|
|
concept_name = 'Stable_Zero'
|
|
subdirs.append(file_output.sanitize_path_part(Path(concept_name).stem.upper()))
|
|
|
|
if add_modelname_to_path == True and 'model' in image_metadata:
|
|
if image_metadata.get('model_name'):
|
|
image_metadata['model'] = image_metadata['model_name']
|
|
subdirs.append(file_output.sanitize_path_part(Path(image_metadata['model']).stem.upper()))
|
|
|
|
if subpath_priority == True and 'preferred' in image_metadata and type(image_metadata['preferred']).__name__ == 'dict' and len(image_metadata['preferred']) > 0 and 'subpath' in image_metadata['preferred']:
|
|
if image_metadata['preferred']['subpath'] is not None and len(image_metadata['preferred']['subpath'].strip()) > 0:
|
|
subpath = image_metadata['preferred']['subpath']
|
|
subdirs.append(file_output.sanitize_path_part(subpath))
|
|
elif subpath_priority == False and subpath is not None and subpath != 'None' and len(subpath.strip()) > 0:
|
|
subdirs.append(file_output.sanitize_path_part(subpath))
|
|
|
|
output_path = file_output.append_subdirs_before_stem(output_path, subdirs)
|
|
output_path = file_output.ensure_output_dir(output_path)
|
|
|
|
file, counter = file_output.build_filename_and_counter(
|
|
output_path=output_path,
|
|
prefix=filename_prefix,
|
|
delimiter=delimiter,
|
|
number_padding=number_padding,
|
|
number_start=filename_number_start,
|
|
extension=extension,
|
|
overwrite_mode=overwrite_mode,
|
|
)
|
|
|
|
results = list()
|
|
# for image in images:
|
|
image = images[0]
|
|
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 png_embed_workflow == True:
|
|
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]))
|
|
|
|
exif_metadata_A11 = None
|
|
try:
|
|
if 'positive' in image_metadata and 'negative' in image_metadata:
|
|
a11samplername = exif_data_checker.comfy_samplers2a11(image_metadata['sampler'], image_metadata['scheduler'])
|
|
image_metadata['vae'] = 'Baked VAE'
|
|
if 'model_hash' not in image_metadata:
|
|
image_metadata['model_hash'] = 'unknown'
|
|
if 'model' in image_metadata:
|
|
checkpointpaths = folder_paths.get_folder_paths("checkpoints")[0]
|
|
model_full_path = checkpointpaths + os.sep + image_metadata['model']
|
|
is_link = os.path.islink(str(model_full_path))
|
|
if is_link == False:
|
|
image_metadata['model_hash'] = exif_data_checker.get_model_hash(model_full_path)
|
|
|
|
if 'is_sdxl' not in image_metadata:
|
|
image_metadata['vae'] = 'Baked VAE'
|
|
else:
|
|
if image_metadata['is_sdxl'] == 1:
|
|
image_metadata['vae'] = image_metadata['vae_name_sdxl']
|
|
else:
|
|
image_metadata['vae'] = image_metadata['vae_name_sd']
|
|
if image_metadata['vae'] is None:
|
|
image_metadata['vae'] = 'Baked VAE'
|
|
|
|
exif_metadata_A11 = f"""{image_metadata['positive']}
|
|
Negative prompt: {image_metadata['negative']}
|
|
Steps: {str(image_metadata['steps'])}, Sampler: {a11samplername}, CFG scale: {str(image_metadata['cfg'])}, Seed: {str(image_metadata['seed'])}, Size: {str(image_metadata['width'])}x{str(image_metadata['height'])}, Model hash: {image_metadata['model_hash']}, Model: {Path((image_metadata['model'])).stem}, VAE: {image_metadata['vae']}"""
|
|
except Exception:
|
|
print('Cannot save A1111 compatible data')
|
|
|
|
try:
|
|
output_file = os.path.abspath(os.path.join(output_path, file))
|
|
exif_metadata_json = image_metadata
|
|
if extension == 'png':
|
|
if png_embed_data == True:
|
|
metadata.add_text("gendata", json.dumps(exif_metadata_json))
|
|
print(f"{extension} Image file saved with description info: {output_file}")
|
|
#img.save(output_file, pnginfo=metadata, optimize=True)
|
|
|
|
if a1111_civitai_meta == True:
|
|
if exif_metadata_A11:
|
|
metadata.add_text("parameters", exif_metadata_A11)
|
|
print(f"{extension} Image file saved with A1111 info: {output_file}")
|
|
img.save(output_file, pnginfo=metadata, optimize=True)
|
|
print(f"{extension} Image file saved with exif: {output_file}")
|
|
|
|
elif extension == 'webp':
|
|
img.save(output_file, quality=quality, exif=metadata)
|
|
print(f"{extension} Image file saved with exif: {output_file}")
|
|
else:
|
|
img.save(output_file, quality=quality, optimize=True)
|
|
metadata = pyexiv2.Image(output_file)
|
|
if image_embed_exif == True:
|
|
metadata.modify_exif({'Exif.Image.ImageDescription': json.dumps(exif_metadata_json)})
|
|
print(f"{extension} Image file saved with description exif: {output_file}")
|
|
if a1111_civitai_meta == True and exif_metadata_A11:
|
|
metadata.modify_exif({'Exif.Photo.UserComment': 'charset=Unicode ' + exif_metadata_A11})
|
|
print(f"{extension} Image file saved with A1111 exif: {output_file}")
|
|
|
|
if a1111_civitai_meta == False and image_embed_exif == False:
|
|
if extension == 'webp':
|
|
img.save(output_file, quality=quality, exif=metadata)
|
|
print(f"{extension} Image file saved without exif: {output_file}")
|
|
else:
|
|
img.save(output_file, quality=quality, optimize=True)
|
|
print(f"{extension} Image file saved without exif: {output_file}")
|
|
|
|
if save_meta_to_json:
|
|
jsonfile = os.path.splitext(output_file)[0] + '.json'
|
|
with open(jsonfile, 'w', encoding='utf-8') as jf:
|
|
json.dump(exif_metadata_json, jf, ensure_ascii=False, indent=4)
|
|
print(f"JSON file saved with generation data: {jsonfile}")
|
|
|
|
except OSError as e:
|
|
print(f'Unable to save file to: {output_file}')
|
|
print(e)
|
|
except Exception as e:
|
|
print('Unable to save file due to the to the following error:')
|
|
print(e)
|
|
|
|
if overwrite_mode == 'false':
|
|
counter += 1
|
|
|
|
filtered_paths = []
|
|
|
|
if filtered_paths:
|
|
for image_path in filtered_paths:
|
|
subfolder = self.get_subfolder_path(image_path, self.output_dir)
|
|
image_data = {
|
|
"filename": os.path.basename(image_path),
|
|
"subfolder": subfolder,
|
|
"type": self.type
|
|
}
|
|
results.append(image_data)
|
|
|
|
metastring = ""
|
|
if image_metadata is not None and len(image_metadata) > 0:
|
|
for key, val in image_metadata.items():
|
|
if len(str(val).strip( '"')) > 0:
|
|
metastring = metastring + ':: ' + key.upper() + ': ' + str(val).strip( '"') + '\n'
|
|
|
|
saved_info = f""":: Time to save: {datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S')}
|
|
:: Output file: {output_file}
|
|
|
|
:: PROCESS INFO ::
|
|
------------------
|
|
{metastring}"""
|
|
|
|
if save_info_to_txt:
|
|
infofile = os.path.splitext(output_file)[0] + '.txt'
|
|
with open(infofile, 'w', encoding='utf-8', newline="") as inf:
|
|
inf.write(saved_info)
|
|
print(f"TXT file saved with generation data: {infofile}")
|
|
|
|
return saved_info, {"ui": {"images": []}}
|
|
|
|
def get_subfolder_path(self, image_path, output_path):
|
|
output_parts = output_path.strip(os.sep).split(os.sep)
|
|
image_parts = image_path.strip(os.sep).split(os.sep)
|
|
common_parts = os.path.commonprefix([output_parts, image_parts])
|
|
subfolder_parts = image_parts[len(common_parts):]
|
|
subfolder_path = os.sep.join(subfolder_parts[:-1])
|
|
return subfolder_path
|
|
|
|
class AnyType(str):
|
|
def __ne__(self, __value: object) -> bool:
|
|
return False
|
|
|
|
any = AnyType("*")
|
|
|
|
class PrimereAnyOutput:
|
|
RETURN_TYPES = ()
|
|
FUNCTION = "show_output"
|
|
OUTPUT_NODE = True
|
|
CATEGORY = TREE_OUTPUTS
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"input": (any, {}),
|
|
},
|
|
}
|
|
|
|
def show_output(self, input = None):
|
|
value = 'None'
|
|
if input is not None:
|
|
try:
|
|
value = json.dumps(input, indent=4)
|
|
except Exception:
|
|
try:
|
|
value = str(input)
|
|
except Exception:
|
|
value = 'Input data exists, but could not be serialized.'
|
|
|
|
return {"ui": {"text": (value.strip( '"'),)}}
|
|
|
|
class PrimereTextOutput:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"text": ("STRING", {"forceInput": True}),
|
|
},
|
|
}
|
|
|
|
INPUT_IS_LIST = True
|
|
RETURN_TYPES = ()
|
|
FUNCTION = "notify"
|
|
OUTPUT_NODE = True
|
|
OUTPUT_IS_LIST = (True,)
|
|
CATEGORY = TREE_OUTPUTS
|
|
|
|
def notify(self, text):
|
|
return {"ui": {"text": text}}
|
|
|
|
class PrimereMetaCollector:
|
|
CATEGORY = TREE_OUTPUTS
|
|
RETURN_TYPES = ("TUPLE",)
|
|
RETURN_NAMES = ("METADATA",)
|
|
FUNCTION = "load_process_meta"
|
|
|
|
INPUT_DICT = {
|
|
"required": {
|
|
"positive": ('STRING', {"forceInput": True, "default": "Red sportcar racing on the street of metropolis"}),
|
|
"negative": ('STRING', {"forceInput": True, "default": "Cute cat, nsfw, nude, nudity, porn"})
|
|
}, "optional": {
|
|
"positive_decoded": ('STRING', {"forceInput": True, "default": "Red sportcar racing on the street of metropolis"}),
|
|
"negative_decoded": ('STRING', {"forceInput": True, "default": "Cute cat, nsfw, nude, nudity, porn"}),
|
|
"control_data": ("TUPLE", {"default": None, "forceInput": True}),
|
|
"seed": ('INT', {"forceInput": True, "default": 1}),
|
|
"t5_xxl_prompt": ('STRING', {"forceInput": True}),
|
|
"positive_l": ('STRING', {"forceInput": True}),
|
|
"negative_l": ('STRING', {"forceInput": True}),
|
|
"width": ('INT', {"forceInput": True, "default": 512}),
|
|
"height": ('INT', {"forceInput": True, "default": 512}),
|
|
"preferred": ("TUPLE", {"default": None, "forceInput": True}),
|
|
"aesthetic_score": ('INT', {"forceInput": True, "default": 0})
|
|
},
|
|
}
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return cls.INPUT_DICT
|
|
|
|
def load_process_meta(self, *args, **kwargs):
|
|
data_json = {}
|
|
|
|
for key, value in self.INPUT_DICT.items():
|
|
for key_l2, value_l2 in value.items():
|
|
if 'default' in value_l2[1]:
|
|
default_value = value_l2[1]['default']
|
|
else:
|
|
default_value = None
|
|
|
|
if key_l2 not in kwargs:
|
|
data_json[key_l2] = default_value
|
|
else:
|
|
data_json[key_l2] = kwargs[key_l2]
|
|
|
|
nested_control = data_json.pop('control_data', None)
|
|
if nested_control and isinstance(nested_control, dict):
|
|
data_json.update(nested_control)
|
|
|
|
return (data_json,)
|
|
|
|
class PrimereKSampler:
|
|
CATEGORY = TREE_OUTPUTS
|
|
RETURN_TYPES = ("LATENT", "TUPLE")
|
|
RETURN_NAMES = ("LATENT", "CONTROL_DATA")
|
|
FUNCTION = "pk_sampler"
|
|
|
|
def __init__(self):
|
|
self.state_hash = False
|
|
self.count = 0
|
|
self.noise_base = 0
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"model": ("MODEL", {"forceInput": True}),
|
|
"seed": ("INT", {"default": 0, "min": 0, "max": utility.MAX_SEED, "forceInput": True}),
|
|
"steps": ("INT", {"default": 20, "min": 1, "max": 10000, "forceInput": True}),
|
|
"cfg": ("FLOAT", {"default": 8.0, "min": 0.0, "max": 100.0, "forceInput": True}),
|
|
"sampler_name": (comfy.samplers.KSampler.SAMPLERS, {"forceInput": True}),
|
|
"scheduler_name": (comfy.samplers.KSampler.SCHEDULERS, {"forceInput": True}),
|
|
"positive": ("CONDITIONING", ),
|
|
"negative": ("CONDITIONING", ),
|
|
"latent_image": ("LATENT", ),
|
|
"denoise": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01}),
|
|
"variation_extender": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 1.0, "step": 0.01}),
|
|
"variation_batch_step": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 0.5, "step": 0.01}),
|
|
"variation_level": ("BOOLEAN", {"default": False, "label_on": "Maximize", "label_off": "Off"}),
|
|
"device": (["DEFAULT", "GPU", "CPU"], {"default": 'DEFAULT'})
|
|
},
|
|
"optional": {
|
|
"model_concept": ("STRING", {"default": "Auto", "forceInput": True}),
|
|
"control_data": ("TUPLE", {"default": None}),
|
|
},
|
|
"hidden": {
|
|
"extra_pnginfo": "EXTRA_PNGINFO",
|
|
"prompt": "PROMPT"
|
|
}
|
|
}
|
|
|
|
@classmethod
|
|
def IS_CHANGED(self, **kwargs):
|
|
if kwargs['variation_extender'] > 0 or kwargs['device'] != 'DEFAULT' or kwargs['variation_batch_step'] > 0 or kwargs['variation_level'] == True:
|
|
return float("NaN")
|
|
|
|
def pk_sampler(self, model, seed, steps, cfg, sampler_name, scheduler_name, positive, negative, latent_image, extra_pnginfo, prompt, model_concept = "Auto", control_data = None, denoise=1.0, variation_extender = 0, variation_batch_step = 0, variation_level = False, model_sampling = 2.5, device = 'DEFAULT', align_your_steps = False):
|
|
timestamp_start = time.time()
|
|
if control_data is not None:
|
|
align_your_steps = control_data.get('align_your_steps', align_your_steps)
|
|
model_sampling = control_data.get('model_sampling', model_sampling)
|
|
if control_data is not None and len(control_data) > 0 and 'exif_status' in control_data and control_data['exif_status'] == 'SUCCEED':
|
|
if 'sampler_settings' in control_data and len(control_data['sampler_settings']) > 0 and 'setup_states' in control_data and 'sampler_setup' in control_data['setup_states']:
|
|
if control_data['setup_states']['sampler_setup'] == True:
|
|
variation_batch_step = 0
|
|
variation_level = False
|
|
denoise = control_data['sampler_settings']['denoise']
|
|
device = control_data['sampler_settings']['device']
|
|
align_your_steps = control_data['sampler_settings']['align_your_steps']
|
|
model_sampling = control_data['sampler_settings']['model_sampling']
|
|
if control_data['sampler_settings']['variation_level'] == True:
|
|
variation_extender = control_data['sampler_settings']['noise_constant']
|
|
else:
|
|
variation_extender = control_data['sampler_settings']['variation_extender_original']
|
|
|
|
samples_out = latent_image
|
|
# out = latent_image.copy()
|
|
variation_extender_original = variation_extender
|
|
variation_batch_step_original = variation_batch_step
|
|
variation_limit = 0.12
|
|
|
|
def check_state(self, extra_pnginfo, prompt):
|
|
old = self.state_hash
|
|
self.state_hash = utility.collect_state(extra_pnginfo, prompt)
|
|
if self.state_hash == old:
|
|
self.count += 1
|
|
return self.count
|
|
self.count = 0
|
|
return self.count
|
|
batch_counter = int(check_state(self, extra_pnginfo, prompt)) + 1
|
|
|
|
def new_state():
|
|
random.seed(datetime.datetime.now().timestamp())
|
|
return random.randint(1000, utility.MAX_SEED)
|
|
|
|
def get_noise_extender(variation_limit, state_random):
|
|
random.seed(state_random)
|
|
noise_extender_low = round(random.uniform(0.00, variation_limit), 2)
|
|
noise_extender_high = round(random.uniform((1 - variation_limit), 1), 2)
|
|
noise_extender = random.choice([noise_extender_low, noise_extender_high])
|
|
return noise_extender
|
|
|
|
state_random = int(new_state())
|
|
noise_extender_ksampler = get_noise_extender(variation_limit, state_random)
|
|
|
|
random.seed(state_random)
|
|
noise_extender_cascade = round(random.uniform((1 - (variation_limit + 0.4)), 1), 2)
|
|
|
|
if variation_batch_step_original > 0:
|
|
if batch_counter > 0:
|
|
variation_batch_step = variation_batch_step_original * batch_counter
|
|
variation_extender = round(variation_extender_original + variation_batch_step, 2)
|
|
noise_extender_ksampler = variation_extender
|
|
noise_extender_cascade = variation_extender
|
|
|
|
elif variation_batch_step_original == 0 and variation_extender_original > 0:
|
|
variation_extender = variation_extender_original
|
|
if batch_counter > 1:
|
|
variation_extender = variation_extender_original + (batch_counter / 100)
|
|
noise_extender_ksampler = variation_extender
|
|
noise_extender_cascade = variation_extender
|
|
|
|
if variation_extender > 1:
|
|
random.seed(batch_counter)
|
|
variation_extender = round(random.uniform((1 - variation_limit), 1), 2)
|
|
noise_extender_ksampler = variation_extender
|
|
noise_extender_cascade = variation_extender
|
|
|
|
noise_constant = noise_extender_ksampler
|
|
WORKFLOWDATA = extra_pnginfo['workflow']['nodes']
|
|
|
|
refiner_model_data = None
|
|
if isinstance(model, dict) and 'main' in model:
|
|
refiner_model_data = model.get('refiner')
|
|
model = model['main']
|
|
refiner_cond_pos = None
|
|
refiner_cond_neg = None
|
|
if isinstance(positive, dict) and 'main' in positive:
|
|
refiner_cond_pos = positive.get('refiner')
|
|
positive = positive['main']
|
|
if isinstance(negative, dict) and 'main' in negative:
|
|
refiner_cond_neg = negative.get('refiner')
|
|
negative = negative['main']
|
|
|
|
match model_concept:
|
|
case 'SANA1024' | 'SANA512':
|
|
if scheduler_name == 'flow_dpm-solver':
|
|
device = model['device']
|
|
samples_out = primeresamplers.PSamplerSana(self, device, seed, model,
|
|
steps, cfg, sampler_name, scheduler_name,
|
|
positive, negative,
|
|
latent_image, denoise,
|
|
variation_extender, variation_batch_step_original,
|
|
batch_counter, variation_extender_original,
|
|
variation_batch_step, variation_level, variation_limit,
|
|
align_your_steps, noise_extender_ksampler, WORKFLOWDATA, prompt)[0]
|
|
else:
|
|
if device == 'DEFAULT':
|
|
device = model_management.get_torch_device()
|
|
latentWidth, latentHeigth = utility.getLatentSize(latent_image)
|
|
|
|
latent = torch.zeros([1, 32, (latentHeigth * 8) // 32, (latentWidth * 8) // 32], device=device)
|
|
latent_image = {"samples": latent}
|
|
samples_out = primeresamplers.PKSampler(self, device, seed, model,
|
|
steps, cfg, sampler_name, scheduler_name,
|
|
positive, negative,
|
|
latent_image, denoise,
|
|
variation_extender, variation_batch_step_original, batch_counter, variation_extender_original, variation_batch_step, variation_level, variation_limit,
|
|
align_your_steps, noise_extender_ksampler, None)[0]
|
|
|
|
try:
|
|
comfy.model_management.soft_empty_cache()
|
|
comfy.model_management.cleanup_models(True)
|
|
except Exception:
|
|
print('No need to clear cache...')
|
|
|
|
case "PixartSigma":
|
|
samples_out = primeresamplers.PSamplerPixart(self, device, seed, model,
|
|
steps, cfg, sampler_name, scheduler_name,
|
|
positive, negative,
|
|
latent_image, denoise,
|
|
variation_extender, variation_batch_step_original,
|
|
batch_counter, variation_extender_original,
|
|
variation_batch_step, variation_level, variation_limit,
|
|
align_your_steps, noise_extender_ksampler, control_data)[0]
|
|
case "KwaiKolors":
|
|
samples_out = primeresamplers.PSamplerKOROLS(self, model, seed, cfg, positive, negative, latent_image, steps, denoise, sampler_name, scheduler_name, model_sampling, 1000)[0]
|
|
case "SD3":
|
|
samples_out = primeresamplers.PSamplerSD3(self, model, seed, cfg, positive, negative, latent_image, steps, denoise, sampler_name, scheduler_name, model_sampling, 1000)[0]
|
|
case "Turbo":
|
|
samples_out = primeresamplers.PTurboSampler(model, seed, cfg, positive, negative, latent_image, steps, denoise, sampler_name)[0]
|
|
case "StableCascade":
|
|
align_your_steps = False
|
|
noise_constant = noise_extender_cascade
|
|
samples_out = primeresamplers.PCascadeSampler(self, model, seed, steps, cfg, sampler_name, scheduler_name, positive, negative, latent_image, denoise, device, variation_level, variation_limit, variation_extender_original, variation_batch_step_original, variation_extender, variation_batch_step, batch_counter, noise_extender_cascade)[0]
|
|
case "Hyper":
|
|
CONCEPT_SELECTOR = control_data.get('model_concept') if control_data else None
|
|
OriginalBaseModel = control_data.get('model_name') if control_data else None
|
|
fullpathFile = folder_paths.get_full_path('checkpoints', OriginalBaseModel)
|
|
is_link = os.path.islink(str(fullpathFile))
|
|
HYPERSD_SELECTOR = 'UNET' if is_link else 'LORA'
|
|
HYPERSD_SAMPLER = True
|
|
|
|
if model_concept == 'Hyper' and (CONCEPT_SELECTOR == 'Hyper' or CONCEPT_SELECTOR == 'Auto') and steps == 12 and HYPERSD_SELECTOR == 'LORA' and HYPERSD_SAMPLER == True:
|
|
cfg = float(control_data.get('cfg', 3.80))
|
|
scheduler_name = control_data.get('scheduler_name', "normal")
|
|
samples_out = primeresamplers.PSamplerHyper(self, extra_pnginfo, model, seed, steps, cfg, positive, negative, sampler_name, scheduler_name, latent_image, denoise, prompt, control_data)[0]
|
|
|
|
case 'QwenGen' | 'QwenEdit':
|
|
align_your_steps = False
|
|
samples_out = primeresamplers.PKSampler(self, device, seed, model,
|
|
steps, cfg, sampler_name, scheduler_name,
|
|
positive, negative,
|
|
latent_image, denoise,
|
|
0, variation_batch_step_original, batch_counter, variation_extender_original, 0, False, variation_limit,
|
|
align_your_steps, noise_extender_ksampler, None)[0]
|
|
|
|
case 'Chroma':
|
|
align_your_steps = False
|
|
samples_out = primeresamplers.PSamplerChroma(self, model, seed, cfg, positive, negative, scheduler_name, sampler_name, steps, denoise, latent_image)[0]
|
|
|
|
case 'Z-Image':
|
|
align_your_steps = False
|
|
model = nodes_model_advanced.ModelSamplingSD3.patch(self, model, 2.8, 1.0)[0]
|
|
negative = nodes.ConditioningZeroOut.zero_out(self, negative)[0]
|
|
samples_out = primeresamplers.PKSampler(self, device, seed, model,
|
|
steps, cfg, sampler_name, scheduler_name,
|
|
positive, negative,
|
|
latent_image, denoise,
|
|
variation_extender, variation_batch_step_original, batch_counter, variation_extender_original, variation_batch_step, variation_level, variation_limit,
|
|
align_your_steps, noise_extender_ksampler, None)[0]
|
|
|
|
case 'Flux':
|
|
FLUX_SAMPLER = control_data.get('sampler', 'ksampler') if control_data else 'ksampler'
|
|
FLUX_GUIDANCE = float(control_data.get('guidance', 3.5)) if control_data else 3.5
|
|
align_your_steps = False
|
|
if FLUX_SAMPLER == 'custom_advanced':
|
|
samples_out = primeresamplers.PSamplerAdvanced(self, model, seed, FLUX_GUIDANCE, positive, scheduler_name, sampler_name, steps, denoise, latent_image)[0]
|
|
elif FLUX_SAMPLER == 'ksampler':
|
|
CONDITIONING_POS = nodes_flux.FluxGuidance.execute(positive, FLUX_GUIDANCE)[0] if FLUX_GUIDANCE > 0 else positive
|
|
if control_data is not None and float(control_data.get('cfg', 2.0)) < 1.2:
|
|
CONDITIONING_NEG = CONDITIONING_POS
|
|
else:
|
|
CONDITIONING_NEG = nodes_flux.FluxGuidance.execute(negative, FLUX_GUIDANCE)[0]
|
|
samples_out = primeresamplers.PKSampler(self, device, seed, model,
|
|
steps, cfg, sampler_name, scheduler_name,
|
|
CONDITIONING_POS, CONDITIONING_NEG,
|
|
latent_image, denoise,
|
|
variation_extender, variation_batch_step_original, batch_counter, variation_extender_original, variation_batch_step, variation_level, variation_limit,
|
|
align_your_steps, noise_extender_ksampler, None, control_data)[0]
|
|
else:
|
|
samples_out = primeresamplers.PKSampler(self, device, seed, model,
|
|
steps, cfg, sampler_name, scheduler_name,
|
|
positive, negative,
|
|
latent_image, denoise,
|
|
variation_extender, variation_batch_step_original, batch_counter, variation_extender_original, variation_batch_step, variation_level, variation_limit,
|
|
align_your_steps, noise_extender_ksampler, None, control_data)[0]
|
|
|
|
case _:
|
|
if model_concept == 'AuraFlow' and model_sampling is not None and model_sampling > 0:
|
|
model = nodes_model_advanced.ModelSamplingSD3.patch(self, model, model_sampling, 1.0)[0]
|
|
samples_out = primeresamplers.PKSampler(self, device, seed, model,
|
|
steps, cfg, sampler_name, scheduler_name,
|
|
positive, negative,
|
|
latent_image, denoise,
|
|
variation_extender, variation_batch_step_original, batch_counter, variation_extender_original, variation_batch_step, variation_level, variation_limit,
|
|
align_your_steps, noise_extender_ksampler, None, control_data)[0]
|
|
|
|
if refiner_model_data is not None:
|
|
samples_out = primeresamplers._run_refiner_pass(self, refiner_model_data, refiner_cond_pos, refiner_cond_neg, samples_out, control_data, seed)
|
|
|
|
if control_data is not None:
|
|
control_data['sampler_settings'] = {}
|
|
control_data['sampler_settings']['denoise'] = denoise
|
|
control_data['sampler_settings']['variation_extender_original'] = variation_extender_original
|
|
control_data['sampler_settings']['variation_batch_step_original'] = variation_batch_step_original
|
|
control_data['sampler_settings']['variation_level'] = variation_level
|
|
control_data['sampler_settings']['device'] = device
|
|
control_data['sampler_settings']['align_your_steps'] = align_your_steps
|
|
control_data['sampler_settings']['noise_constant'] = noise_constant
|
|
control_data['sampler_settings']['variation_seed'] = seed
|
|
control_data['sampler_settings']['batch_counter'] = batch_counter
|
|
control_data['sampler_settings']['model_sampling'] = model_sampling
|
|
|
|
timestamp_diff = int(time.time() - timestamp_start)
|
|
is_random_model = utility.getDataFromWorkflowByName(WORKFLOWDATA, 'PrimereVisualCKPT', 'random_model', prompt)
|
|
selected_model = control_data.get('model_name') if control_data else utility.getDataFromWorkflowByName(WORKFLOWDATA, 'PrimereVisualCKPT', 'base_model', prompt)
|
|
if is_random_model == True:
|
|
fullSource = PrimereVisualCKPT.allModels
|
|
slashIndex = selected_model.find('\\')
|
|
if slashIndex > 0:
|
|
subdirType = selected_model[0: slashIndex] + '\\'
|
|
models_by_path = list(filter(lambda x: x.startswith(subdirType), fullSource))
|
|
random.seed(seed)
|
|
selected_model = random.choice(models_by_path)
|
|
|
|
if selected_model is not None:
|
|
modelname_only = Path(selected_model).stem
|
|
model_samplingtime = utility.get_value_from_cache('model_samplingtime', modelname_only)
|
|
if model_samplingtime is None:
|
|
utility.add_value_to_cache('model_samplingtime', modelname_only, '1|' + str(timestamp_diff))
|
|
else:
|
|
model_samplingtime_list = model_samplingtime.split("|")
|
|
counter = str(int(model_samplingtime_list[0]) + 1)
|
|
diffvalue = str(int(model_samplingtime_list[1]) + timestamp_diff)
|
|
utility.add_value_to_cache('model_samplingtime', modelname_only, counter + '|' + diffvalue)
|
|
|
|
return (samples_out, control_data)
|
|
|
|
class PrimerePreviewImage():
|
|
CATEGORY = TREE_OUTPUTS
|
|
RETURN_TYPES = ()
|
|
OUTPUT_NODE = True
|
|
FUNCTION = "preview_img_saver"
|
|
|
|
image_path = folder_paths.get_output_directory()
|
|
|
|
def __init__(self):
|
|
self.output_dir = folder_paths.get_temp_directory()
|
|
self.type = "temp"
|
|
self.prefix_append = "_temp_" + ''.join(random.choice("abcdefghijklmnopqrstupvxyz") for x in range(5))
|
|
self.compress_level = 1
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"image_save_as": ("BOOLEAN", {"default": False, "label_on": "Save as preview", "label_off": "Save as any..."}),
|
|
"image_type": (['jpeg', 'png', 'webp'], {"default": "jpeg"}),
|
|
"image_resize": ("INT", {"default": 0, "min": 0, "max": utility.MAX_RESOLUTION, "step": 64}),
|
|
"image_quality": ("INT", {"default": 95, "min": 10, "max": 100, "step": 5}),
|
|
"preview_target": (['Checkpoint', 'CSV Prompt', 'Lora', 'Lycoris', 'Hypernetwork', 'Embedding'],),
|
|
"preview_save_mode": (['Overwrite', 'Keep', 'Join horizontal', 'Join vertical'], {"default": "Overwrite"}),
|
|
"embed_metadata": ("BOOLEAN", {"default": False, "label_on": "Embed metadata", "label_off": "No metadata"}),
|
|
"auto_save_path": ("BOOLEAN", {"default": False, "label_on": "Comfy output folder", "label_off": "Temp folder, will be deleted"}),
|
|
},
|
|
"optional": {
|
|
"images": ("IMAGE", {"default": None}),
|
|
"image_metadata": ('TUPLE', {"forceInput": True}),
|
|
},
|
|
"hidden": {
|
|
"extra_pnginfo": "EXTRA_PNGINFO",
|
|
"prompt": "PROMPT",
|
|
"image_path": (cls.image_path,),
|
|
"id": "UNIQUE_ID",
|
|
},
|
|
}
|
|
|
|
def preview_img_saver(self, image_save_as, image_type, image_resize, image_quality, preview_target, preview_save_mode, embed_metadata, auto_save_path = False, images=None, image_metadata=None, **kwargs):
|
|
if auto_save_path:
|
|
self.output_dir = folder_paths.get_output_directory()
|
|
self.type = "output"
|
|
|
|
if images is None or type(images).__name__ != "Tensor":
|
|
INVALID_IMAGE_PATH = os.path.join(PRIMERE_ROOT, 'front_end', 'images')
|
|
INVALID_IMAGE = os.path.join(INVALID_IMAGE_PATH, "invalid.jpg")
|
|
images = utility.ImageLoaderFromPath(INVALID_IMAGE)
|
|
|
|
VISUAL_NODE_NAMES = ['PrimereVisualCKPT', 'PrimereVisualLORA', 'PrimereVisualEmbedding', 'PrimereVisualHypernetwork', 'PrimereVisualLYCORIS', 'PrimereVisualStyle']
|
|
VISUAL_NODE_FILENAMES = ['PrimereVisualCKPT', 'PrimereVisualLORA', 'PrimereVisualEmbedding', 'PrimereVisualHypernetwork', 'PrimereVisualLYCORIS']
|
|
WIDGET_DATA = {
|
|
"PrimereVisualCKPT": ['base_model'],
|
|
"PrimereVisualStyle": ['styles'],
|
|
"PrimereVisualLORA": ['lora_1', 'lora_2', 'lora_3', 'lora_4', 'lora_5', 'lora_6'],
|
|
"PrimereVisualEmbedding": ['embedding_1', 'embedding_2', 'embedding_3', 'embedding_4', 'embedding_5', 'embedding_6'],
|
|
"PrimereVisualHypernetwork": ['hypernetwork_1', 'hypernetwork_2', 'hypernetwork_3', 'hypernetwork_4', 'hypernetwork_5', 'hypernetwork_6'],
|
|
"PrimereVisualLYCORIS": ['lycoris_1', 'lycoris_2', 'lycoris_3', 'lycoris_4', 'lycoris_5', 'lycoris_6'],
|
|
}
|
|
|
|
WORKFLOWDATA = kwargs['extra_pnginfo']['workflow']['nodes']
|
|
VISUAL_DATA = {}
|
|
for NODE_ITEMS in WORKFLOWDATA:
|
|
ITEM_TYPE = NODE_ITEMS['type']
|
|
if ITEM_TYPE in VISUAL_NODE_NAMES and ITEM_TYPE in WIDGET_DATA:
|
|
REQUIRED_DATA_NAMES = WIDGET_DATA[ITEM_TYPE]
|
|
if len(REQUIRED_DATA_NAMES) > 0:
|
|
VALUE_LIST = []
|
|
VALUE_LIST_ORIGINAL = []
|
|
WIDGET_STATE = True
|
|
for DATA_NAME in REQUIRED_DATA_NAMES:
|
|
if type(DATA_NAME).__name__ == 'str':
|
|
WIDGET_VALUE_ORIGINAL = utility.getDataFromWorkflowByName(WORKFLOWDATA, ITEM_TYPE, DATA_NAME, kwargs['prompt'])
|
|
if DATA_NAME[-1].isdigit():
|
|
USE_WIDGET_NAME = 'use_' + DATA_NAME
|
|
WIDGET_STATE = utility.getDataFromWorkflowByName(WORKFLOWDATA, ITEM_TYPE, USE_WIDGET_NAME, kwargs['prompt'])
|
|
|
|
if ITEM_TYPE in VISUAL_NODE_FILENAMES:
|
|
REPLACED_WIDGETS = Path(WIDGET_VALUE_ORIGINAL).stem.replace(' ', '_')
|
|
else:
|
|
REPLACED_WIDGETS = WIDGET_VALUE_ORIGINAL.replace(' ', '_')
|
|
|
|
if WIDGET_STATE == True and REPLACED_WIDGETS not in VALUE_LIST:
|
|
VALUE_LIST.append(REPLACED_WIDGETS)
|
|
VALUE_LIST_ORIGINAL.append(WIDGET_VALUE_ORIGINAL)
|
|
|
|
VISUAL_DATA[ITEM_TYPE] = VALUE_LIST
|
|
VISUAL_DATA[ITEM_TYPE + '_ORIGINAL'] = VALUE_LIST_ORIGINAL
|
|
|
|
extra_pnginfo_embed = None
|
|
prompt_embed = None
|
|
if embed_metadata and not args.disable_metadata:
|
|
prompt_embed = kwargs.get('prompt')
|
|
if image_metadata is not None and isinstance(image_metadata, dict):
|
|
try:
|
|
extra_pnginfo_embed = {"gendata": image_metadata}
|
|
except Exception:
|
|
pass
|
|
elif image_metadata is None and kwargs.get('extra_pnginfo') is not None:
|
|
extra_pnginfo_embed = kwargs.get('extra_pnginfo')
|
|
|
|
if auto_save_path:
|
|
results = nodes.SaveImage.save_images(self, images, filename_prefix="Primere_ComfyUI", prompt=prompt_embed, extra_pnginfo=extra_pnginfo_embed)
|
|
else:
|
|
r1 = random.randint(1000, 9999)
|
|
temp_filename = f"Primere_ComfyUI_{r1}.png"
|
|
os.makedirs(folder_paths.get_temp_directory(), exist_ok=True)
|
|
TEMP_FILE = os.path.join(folder_paths.get_temp_directory(), temp_filename)
|
|
utility.tensor_to_image(images[0]).save(TEMP_FILE)
|
|
results = {"ui": {"images": [{"filename": temp_filename, "subfolder": "", "type": "temp"}]}, "result": ()}
|
|
|
|
ui_images = results.get('ui', {}).get('images', [])
|
|
VISUAL_DATA['SaveImages'] = ui_images
|
|
VISUAL_DATA['node_id'] = kwargs.get('id')
|
|
VISUAL_DATA['ImagePath'] = folder_paths.get_output_directory() if auto_save_path else folder_paths.get_temp_directory()
|
|
PromptServer.instance.send_sync("getVisualTargets", VISUAL_DATA)
|
|
|
|
return results
|
|
|
|
class PrimereAestheticCKPTScorer:
|
|
CATEGORY = TREE_OUTPUTS
|
|
RETURN_TYPES = ("INT",)
|
|
RETURN_NAMES = ("SCORE",)
|
|
OUTPUT_NODE = True
|
|
FUNCTION = "aesthetic_scorer"
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"get_aesthetic_score": ("BOOLEAN", {"default": False}),
|
|
"add_to_checkpoint": ("BOOLEAN", {"default": False}),
|
|
"add_to_saved_prompt": ("BOOLEAN", {"default": False}),
|
|
"image": ("IMAGE", ),
|
|
},
|
|
"optional": {
|
|
"control_data": ('TUPLE', {"forceInput": True}),
|
|
},
|
|
"hidden": {
|
|
"extra_pnginfo": "EXTRA_PNGINFO",
|
|
"prompt": "PROMPT"
|
|
},
|
|
}
|
|
|
|
def aesthetic_scorer(self, image, get_aesthetic_score, add_to_checkpoint, add_to_saved_prompt, prompt, dual_mode = True, control_data = None, **kwargs):
|
|
final_prediction = '*** Aesthetic scorer off ***'
|
|
models = []
|
|
AE_SCORE_MIN = None
|
|
AE_SCORE_MAX = None
|
|
if 'extra_pnginfo' in kwargs:
|
|
WORKFLOWDATA = kwargs['extra_pnginfo']['workflow']['nodes']
|
|
AE_SCORE_MIN = utility.getDataFromWorkflowByName(WORKFLOWDATA, 'PrimereVisualCKPT', 'aescore_percent_min', prompt)
|
|
AE_SCORE_MAX = utility.getDataFromWorkflowByName(WORKFLOWDATA, 'PrimereVisualCKPT', 'aescore_percent_max', prompt)
|
|
|
|
def pipe(model):
|
|
return pipeline(task="image-classification", model=model, device=model_management.get_torch_device())
|
|
|
|
if (get_aesthetic_score == True):
|
|
AE_MODEL_ROOT = os.path.join(folder_paths.models_dir, 'aesthetic')
|
|
AEMODELS_ENCODERS_PATHS = utility.getValidAscorerPaths(AE_MODEL_ROOT)
|
|
if len(AEMODELS_ENCODERS_PATHS) > 0:
|
|
if 'cafe_aesthetic' not in AEMODELS_ENCODERS_PATHS:
|
|
final_prediction = '*** Missing aesthetic model ***'
|
|
if 'cafe_style' not in AEMODELS_ENCODERS_PATHS:
|
|
final_prediction = '*** Missing style model ***'
|
|
if 'cafe_style' in AEMODELS_ENCODERS_PATHS and 'cafe_aesthetic' in AEMODELS_ENCODERS_PATHS:
|
|
ae_model_access = os.path.join(AE_MODEL_ROOT, 'cafe_aesthetic')
|
|
style_model_access = os.path.join(AE_MODEL_ROOT, 'cafe_style')
|
|
if os.path.isdir(ae_model_access) == True and os.path.isdir(style_model_access) == True:
|
|
if dual_mode == True:
|
|
models.append({"pipe": pipe(ae_model_access), "weights": [0.0, 1.0], })
|
|
models.append({"pipe": pipe(style_model_access), "weights": [1.0, 0.75, 0.5, 0.0, 0.0], })
|
|
final_divider = 2
|
|
else:
|
|
models.append({"pipe": pipe(ae_model_access), "weights": [0.0, 1.0], })
|
|
final_divider = 1
|
|
try:
|
|
count = 1
|
|
pil_images = image.permute(0, 3, 1, 2)
|
|
pil_images = torch.clamp(pil_images * 255, 0, 255)
|
|
pil_images = pil_images.to("cpu", torch.uint8)
|
|
pil_images = [TF.to_pil_image(i) for i in pil_images]
|
|
scores = {i: 0.0 for i in range(image.shape[0])}
|
|
for model in models:
|
|
pipe = model["pipe"]
|
|
weights = model["weights"]
|
|
labels = pipe.model.config.id2label
|
|
w_len = len(weights)
|
|
w_sum = sum(weights)
|
|
w_map = {labels[i]: weights[i] for i in range(w_len)}
|
|
values = pipe(pil_images, top_k=w_len)
|
|
for index, value in enumerate(values):
|
|
score = [v["score"] * w_map[v["label"]] for v in value]
|
|
scores[index] += sum(score) / w_sum
|
|
scores = sorted(scores.items(), key=lambda k: k[1], reverse=True)[:count]
|
|
final_score = ", ".join([f"{v:.3f}" for k, v in scores])
|
|
final_prediction = int((float(final_score) * 1000) / final_divider)
|
|
except Exception:
|
|
final_prediction = '*** Invalid input image ***'
|
|
else:
|
|
final_prediction = '*** No aesthetic models downloaded ***'
|
|
else:
|
|
final_prediction = '*** No aesthetic models downloaded ***'
|
|
|
|
if type(final_prediction) != str:
|
|
final_prediction = str(final_prediction)
|
|
|
|
if control_data is not None and final_prediction.isdigit():
|
|
if add_to_checkpoint == True and (control_data['model_concept']):
|
|
if 'model_name' in control_data:
|
|
AE_SCORE_MIN = utility.getDataFromWorkflowByName(WORKFLOWDATA, 'PrimereVisualCKPT', 'aescore_percent_min', prompt)
|
|
AE_SCORE_MAX = utility.getDataFromWorkflowByName(WORKFLOWDATA, 'PrimereVisualCKPT', 'aescore_percent_max', prompt)
|
|
selected_model = control_data['model_name']
|
|
modelname_only = Path(selected_model).stem
|
|
model_ascore = utility.get_value_from_cache('model_ascores', modelname_only)
|
|
if model_ascore is None:
|
|
utility.add_value_to_cache('model_ascores', modelname_only, '1|' + final_prediction)
|
|
else:
|
|
model_ascore_list = model_ascore.split("|")
|
|
counter = str(int(model_ascore_list[0]) + 1)
|
|
score = str(int(model_ascore_list[1]) + int(final_prediction))
|
|
utility.add_value_to_cache('model_ascores', modelname_only, counter + '|' + score)
|
|
|
|
if add_to_saved_prompt == True and final_prediction.isdigit():
|
|
if 'positive' in control_data:
|
|
selectedStyle = utility.getDataFromWorkflowByName(WORKFLOWDATA, 'PrimereVisualStyle', 'styles', prompt)
|
|
if selectedStyle is None:
|
|
selectedStyle = utility.getDataFromWorkflowByName(WORKFLOWDATA, 'PrimereStyleLoader', 'styles', prompt)
|
|
|
|
is_random_style = utility.getDataFromWorkflowByName(WORKFLOWDATA, 'PrimereVisualStyle', 'random_prompt', prompt)
|
|
if is_random_style == True:
|
|
styles_csv = PrimereVisualStyle.styles_csv
|
|
seed = utility.getDataFromWorkflowByName(WORKFLOWDATA, 'PrimereFastSeed', 'seed', prompt)
|
|
random.seed(seed)
|
|
styleKey = styles_csv['name'] == selectedStyle
|
|
try:
|
|
preferred_subpath = styles_csv[styleKey]['preferred_subpath'].values[0]
|
|
except Exception:
|
|
preferred_subpath = ''
|
|
if str(preferred_subpath) == "nan":
|
|
resultsBySubpath = styles_csv[styles_csv['preferred_subpath'].isnull()]
|
|
else:
|
|
resultsBySubpath = styles_csv[styles_csv['preferred_subpath'] == preferred_subpath]
|
|
selectedStyle = random.choice(list(resultsBySubpath['name']))
|
|
|
|
if selectedStyle is not None:
|
|
STYLE_DIR = os.path.join(PRIMERE_ROOT, 'stylecsv')
|
|
STYLE_FILE = os.path.join(STYLE_DIR, "styles.csv")
|
|
try:
|
|
STYLE_FILE_EXAMPLE = os.path.join(STYLE_DIR, "styles.example.csv")
|
|
except Exception:
|
|
STYLE_FILE_EXAMPLE = STYLE_FILE
|
|
|
|
if Path(STYLE_FILE).is_file() == True:
|
|
STYLE_SOURCE = STYLE_FILE
|
|
else:
|
|
STYLE_SOURCE = STYLE_FILE_EXAMPLE
|
|
style_data = utility.load_external_csv(STYLE_SOURCE, 0)
|
|
positive_prompt = style_data[style_data['name'] == selectedStyle]['prompt'].values[0]
|
|
if (positive_prompt is not None):
|
|
if len(positive_prompt) > 100:
|
|
positive_prompt = positive_prompt[:100]
|
|
if positive_prompt in control_data['positive']:
|
|
AE_SCORE_MIN = utility.getDataFromWorkflowByName(WORKFLOWDATA, 'PrimereVisualStyle', 'aescore_percent_min', prompt)
|
|
AE_SCORE_MAX = utility.getDataFromWorkflowByName(WORKFLOWDATA, 'PrimereVisualStyle', 'aescore_percent_max', prompt)
|
|
style_ascore = utility.get_value_from_cache('styles_ascores', selectedStyle)
|
|
if style_ascore is None:
|
|
utility.add_value_to_cache('styles_ascores', selectedStyle, '1|' + final_prediction)
|
|
else:
|
|
style_ascore_list = style_ascore.split("|")
|
|
counter = str(int(style_ascore_list[0]) + 1)
|
|
score = str(int(style_ascore_list[1]) + int(final_prediction))
|
|
utility.add_value_to_cache('styles_ascores', selectedStyle, counter + '|' + score)
|
|
|
|
result_int = int(final_prediction) if final_prediction.isdigit() else 0
|
|
if isinstance(final_prediction, str) and final_prediction.isdigit():
|
|
if AE_SCORE_MIN is not None and AE_SCORE_MAX is not None:
|
|
final_prediction = max(0, min(100, int(((int(final_prediction) - AE_SCORE_MIN) / (AE_SCORE_MAX - AE_SCORE_MIN)) * 100)))
|
|
else:
|
|
final_prediction = int(final_prediction)
|
|
else:
|
|
final_prediction = '*** Aesthetic scorer error ***'
|
|
|
|
return {"ui": {"text": [f'{result_int} / {final_prediction}%']}, "result": (result_int,)}
|
|
|
|
class DebugToFile():
|
|
CATEGORY = TREE_OUTPUTS
|
|
RETURN_TYPES = ("STRING",)
|
|
RETURN_NAMES = ("PROMPT+",)
|
|
FUNCTION = "debug_to_file"
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"input_prompt": ("STRING", {"default": "", "forceInput": True}),
|
|
"enhanced_prompt": ("STRING", {"default": "", "forceInput": True}),
|
|
"seed": ("INT", {"default": 1, "forceInput": True}),
|
|
},
|
|
"hidden": {
|
|
"extra_pnginfo": "EXTRA_PNGINFO",
|
|
"prompt": "PROMPT"
|
|
},
|
|
}
|
|
|
|
def debug_to_file(self, **kwargs):
|
|
WORKFLOWDATA = kwargs['extra_pnginfo']['workflow']['nodes']
|
|
prompt = kwargs['prompt']
|
|
|
|
LLM_NAME = utility.getDataFromWorkflowByName(WORKFLOWDATA, 'PrimereLLMEnhancer', 'llm_model_path', prompt)
|
|
LLM_CONF = utility.getDataFromWorkflowByName(WORKFLOWDATA, 'PrimereLLMEnhancer', 'configurator', prompt)
|
|
|
|
json_dir = os.path.join(PRIMERE_ROOT, 'json')
|
|
json_file = os.path.join(json_dir, 'llm_autotest.json')
|
|
|
|
cacheData = {LLM_NAME: {LLM_CONF: [{"input_prompt": kwargs['input_prompt'], "enhanced_prompt": kwargs['enhanced_prompt'], "seed": kwargs['seed']}]}}
|
|
json_object = json.dumps(cacheData, indent=4)
|
|
ifJsonExist = os.path.isfile(json_file)
|
|
if ifJsonExist == True:
|
|
with open(json_file, 'r') as openfile:
|
|
saved_cache = json.load(openfile)
|
|
if LLM_NAME in saved_cache and LLM_CONF in saved_cache[LLM_NAME]:
|
|
cacheData = [{"input_prompt": kwargs['input_prompt'], "enhanced_prompt": kwargs['enhanced_prompt'], "seed": kwargs['seed']}]
|
|
saved_cache[LLM_NAME][LLM_CONF].append(cacheData)
|
|
elif LLM_NAME in saved_cache and LLM_CONF not in saved_cache[LLM_NAME]:
|
|
cacheData = {LLM_CONF: [{"input_prompt": kwargs['input_prompt'], "enhanced_prompt": kwargs['enhanced_prompt'], "seed": kwargs['seed']}]}
|
|
saved_cache[LLM_NAME].update(cacheData)
|
|
else:
|
|
saved_cache.update(cacheData)
|
|
newJsonObject = json.dumps(saved_cache, indent=4)
|
|
with open(json_file, "w", encoding='utf-8') as outfile:
|
|
outfile.write(newJsonObject)
|
|
else:
|
|
with open(json_file, "w", encoding='utf-8') as outfile:
|
|
outfile.write(json_object)
|
|
|
|
return (LLM_NAME,)
|