ComfyUI/custom_nodes/util_nodes.py

114 lines
3.6 KiB
Python

import hashlib
import os
import numpy as np
import folder_paths
class TextLoader:
def __init__(self, event_dispatcher):
self.event_dispatcher = event_dispatcher
@classmethod
def INPUT_TYPES(s):
prompts_dir = folder_paths.prompt_directory
return {"required":
{"prompt_file": (sorted(os.listdir(prompts_dir)),)},
}
CATEGORY = "utils"
RETURN_TYPES = ("TEXT",)
FUNCTION = "load_text"
def load_text(self, prompt_file):
text_file_path = os.path.join(folder_paths.prompt_directory, prompt_file)
with open(text_file_path, 'r') as f:
text = f.read()
return (text,)
@classmethod
def IS_CHANGED(s, prompt_file):
text_file_path = os.path.join(folder_paths.prompt_directory, prompt_file)
m = hashlib.sha256()
with open(text_file_path, 'rb') as f:
m.update(f.read())
return m.digest().hex()
@classmethod
def VALIDATE_INPUTS(s, prompt_file):
text_file_path = os.path.join(folder_paths.prompt_directory, prompt_file)
if not os.path.exists(text_file_path):
return "Invalid text file: {}".format(text_file_path)
return True
class PrintNode:
def __init__(self, event_dispatcher):
self.event_dispatcher = event_dispatcher
@classmethod
def INPUT_TYPES(s):
return {
"required": {},
"optional": {
"text": ("text",),
"attention": ("ATTENTION",),
"latent": ("LATENT",),
"image": ("IMAGE",),
}
}
@classmethod
def IS_CHANGED(cls, *args, **kwargs):
return True
RETURN_TYPES = ()
FUNCTION = "print_value"
CATEGORY = "utils"
OUTPUT_NODE = True
def print_value(self, text=None, latent=None, attention=None, image=None):
if latent is not None:
latent_hash = hashlib.sha256(latent["samples"].cpu().numpy().tobytes()).hexdigest()
print(f"Latent hash: {latent_hash}")
print(np.array2string(latent["samples"].cpu().numpy(), separator=', '))
output_text = []
# attention[a][b][c][d]
# a: number of steps/sigma in this diffusion process
# b: number of SpatialTransformer or AttentionBlocks used in the middle blocks of the latent diffusion model
# c: number of transformer layers in the SpatialTransformer or AttentionBlocks
# d: attn1, attn2
if attention is not None:
output_text.append(f'attention has {len(attention)} steps\n')
output_text[-1] += f'each step has {len(attention[0])} transformer blocks\n'
output_text[-1] += f'each block has {len(attention[0][0])} transformer layers\n'
output_text[-1] += f'each transformer layer has {len(attention[0][0][0])} attention tensors (attn1, attn2)\n'
output_text[-1] += f'the shape of the attention tensors is {attention[0][0][0][0].shape}\n'
output_text[-1] += f'the first value of the first attention tensor is {attention[0][0][0][0][:1]}\n'
print(output_text[-1])
if text is not None:
output_text.append(text)
print(text)
if image is not None:
_, height, width, _ = image.shape
output_text.append(f"Image dimensions: {width}x{height}\n")
print(output_text[-1])
return {"ui": {"text": "\n".join(output_text)}}
NODE_CLASS_MAPPINGS = {
"PrintNode": PrintNode,
"TextLoader": TextLoader,
}
NODE_DISPLAY_NAME_MAPPINGS = {
"PrintNode": "Print",
"TextLoader": "Text Loader",
}