Robin Collins
d2871d75fd
- Introduced `__init__.py` for package initialization, dependency checking, and node management. - Added `image_nodes.py` with nodes for model name extraction and image saving with metadata. - Created `install.py` for dependency installation and ComfyUI Manager integration setup. - Included configuration files: `node_config.json`, `.manager_info.json`, and `CHANGELOG.md`. - Established `.gitignore` to exclude auto-generated files and unnecessary directories. - Added `README.md` for comprehensive project documentation and usage instructions. - Implemented `DependencyChecker` and `NodeManager` classes for enhanced functionality. - Updated `FILETREE.md` to reflect project structure and organization.
245 lines
8.9 KiB
Python
245 lines
8.9 KiB
Python
import os
|
|
import hashlib
|
|
from datetime import datetime
|
|
import json
|
|
import piexif
|
|
import piexif.helper
|
|
from PIL import Image
|
|
from PIL.PngImagePlugin import PngInfo
|
|
import numpy as np
|
|
import folder_paths
|
|
import comfy.sd
|
|
import comfy.samplers
|
|
from nodes import MAX_RESOLUTION
|
|
|
|
|
|
def parse_name(ckpt_name):
|
|
"""Extract filename without extension from checkpoint path."""
|
|
path = ckpt_name
|
|
filename = path.split("/")[-1]
|
|
filename = filename.split(".")[:-1]
|
|
filename = ".".join(filename)
|
|
return filename
|
|
|
|
|
|
def calculate_sha256(file_path):
|
|
"""Calculate SHA256 hash of a file."""
|
|
sha256_hash = hashlib.sha256()
|
|
|
|
with open(file_path, "rb") as f:
|
|
for byte_block in iter(lambda: f.read(4096), b""):
|
|
sha256_hash.update(byte_block)
|
|
|
|
return sha256_hash.hexdigest()
|
|
|
|
|
|
def handle_whitespace(string: str):
|
|
"""Clean whitespace characters from strings for metadata."""
|
|
return string.strip().replace("\n", " ").replace("\r", " ").replace("\t", " ")
|
|
|
|
|
|
def get_timestamp(time_format):
|
|
"""Generate timestamp string with given format."""
|
|
now = datetime.now()
|
|
try:
|
|
timestamp = now.strftime(time_format)
|
|
except:
|
|
timestamp = now.strftime("%Y-%m-%d-%H%M%S")
|
|
|
|
return timestamp
|
|
|
|
|
|
def make_pathname(filename, seed, modelname, counter, time_format):
|
|
"""Replace placeholders in pathname with actual values."""
|
|
filename = filename.replace("%date", get_timestamp("%Y-%m-%d"))
|
|
filename = filename.replace("%time", get_timestamp(time_format))
|
|
filename = filename.replace("%model", modelname)
|
|
filename = filename.replace("%seed", str(seed))
|
|
filename = filename.replace("%counter", str(counter))
|
|
return filename
|
|
|
|
|
|
def make_filename(filename, seed, modelname, counter, time_format):
|
|
"""Generate final filename with placeholder replacement."""
|
|
filename = make_pathname(filename, seed, modelname, counter, time_format)
|
|
return get_timestamp(time_format) if filename == "" else filename
|
|
|
|
|
|
class ModelNameExtractor:
|
|
"""Extract model names from ComfyUI MODEL objects."""
|
|
|
|
RETURN_TYPES = ("STRING",)
|
|
RETURN_NAMES = ("model_name",)
|
|
FUNCTION = "extract_name"
|
|
CATEGORY = "🖥️🛠TechsToolz"
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"model": ("MODEL",),
|
|
}
|
|
}
|
|
|
|
def extract_name(self, model):
|
|
"""Extract model name from various ComfyUI model types."""
|
|
gguf_filename = getattr(model, "gguf_filename", None)
|
|
if gguf_filename:
|
|
return (gguf_filename,)
|
|
|
|
if hasattr(model, "model") and hasattr(model.model, "model_config"):
|
|
config = model.model.model_config
|
|
if hasattr(config, "unet_config") and hasattr(config.unet_config, "model_path"):
|
|
model_path = config.unet_config.model_path
|
|
if model_path:
|
|
return (os.path.basename(model_path),)
|
|
|
|
if hasattr(model, "model_path"):
|
|
return (os.path.basename(model.model_path),)
|
|
|
|
if hasattr(model, "filename"):
|
|
return (model.filename,)
|
|
|
|
return ("unknown_model",)
|
|
|
|
|
|
class ImageSaveWithMetadata:
|
|
"""Advanced image saving with embedded metadata support."""
|
|
|
|
def __init__(self):
|
|
self.output_dir = folder_paths.output_directory
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"images": ("IMAGE",),
|
|
"filename": ("STRING", {"default": "%time_%seed", "multiline": False}),
|
|
"path": ("STRING", {"default": "", "multiline": False}),
|
|
"extension": (["png", "jpeg", "webp"],),
|
|
"steps": ("INT", {"default": 20, "min": 1, "max": 10000}),
|
|
"cfg": ("FLOAT", {"default": 8.0, "min": 0.0, "max": 100.0}),
|
|
"modelname": ("STRING", {"default": ""}),
|
|
"sampler_name": (comfy.samplers.KSampler.SAMPLERS,),
|
|
"scheduler": (comfy.samplers.KSampler.SCHEDULERS,),
|
|
},
|
|
"optional": {
|
|
"positive": ("STRING", {"default": "unknown", "multiline": True}),
|
|
"negative": ("STRING", {"default": "unknown", "multiline": True}),
|
|
"seed_value": ("INT", {"default": 0, "min": 0, "max": 0xFFFFFFFFFFFFFFFF}),
|
|
"width": ("INT", {"default": 512, "min": 1, "max": MAX_RESOLUTION, "step": 8}),
|
|
"height": ("INT", {"default": 512, "min": 1, "max": MAX_RESOLUTION, "step": 8}),
|
|
"lossless_webp": ("BOOLEAN", {"default": True}),
|
|
"quality_jpeg_or_webp": ("INT", {"default": 100, "min": 1, "max": 100}),
|
|
"counter": ("INT", {"default": 0, "min": 0, "max": 0xFFFFFFFFFFFFFFFF}),
|
|
"time_format": ("STRING", {"default": "%Y-%m-%d-%H%M%S", "multiline": False}),
|
|
},
|
|
"hidden": {"prompt": "PROMPT", "extra_pnginfo": "EXTRA_PNGINFO"},
|
|
}
|
|
|
|
RETURN_TYPES = ()
|
|
FUNCTION = "save_files"
|
|
OUTPUT_NODE = True
|
|
CATEGORY = "🖥️🛠TechsToolz"
|
|
|
|
def save_files(
|
|
self,
|
|
images,
|
|
seed_value,
|
|
steps,
|
|
cfg,
|
|
sampler_name,
|
|
scheduler,
|
|
positive,
|
|
negative,
|
|
modelname,
|
|
quality_jpeg_or_webp,
|
|
lossless_webp,
|
|
width,
|
|
height,
|
|
counter,
|
|
filename,
|
|
path,
|
|
extension,
|
|
time_format,
|
|
prompt=None,
|
|
extra_pnginfo=None,
|
|
):
|
|
"""Save images with embedded metadata."""
|
|
filename = make_filename(filename, seed_value, modelname, counter, time_format)
|
|
path = make_pathname(path, seed_value, modelname, counter, time_format)
|
|
ckpt_path = folder_paths.get_full_path("checkpoints", modelname)
|
|
basemodelname = parse_name(modelname)
|
|
modelhash = calculate_sha256(ckpt_path)[:10]
|
|
comment = f"{handle_whitespace(positive)}\nNegative prompt: {handle_whitespace(negative)}\nSteps: {steps}, Sampler: {sampler_name}{f'_{scheduler}' if scheduler != 'normal' else ''}, CFG Scale: {cfg}, Seed: {seed_value}, Size: {width}x{height}, Model hash: {modelhash}, Model: {basemodelname}, Version: ComfyUI"
|
|
output_path = os.path.join(self.output_dir, path)
|
|
|
|
if output_path.strip() != "":
|
|
if not os.path.exists(output_path.strip()):
|
|
print(f"The path `{output_path.strip()}` specified doesn't exist! Creating directory.")
|
|
os.makedirs(output_path, exist_ok=True)
|
|
|
|
filenames = self.save_images(
|
|
images, output_path, filename, comment, extension, quality_jpeg_or_webp, lossless_webp, prompt, extra_pnginfo
|
|
)
|
|
|
|
subfolder = os.path.normpath(path)
|
|
return {
|
|
"ui": {
|
|
"images": [
|
|
{"filename": filename, "subfolder": subfolder if subfolder != "." else "", "type": "output"} for filename in filenames
|
|
]
|
|
}
|
|
}
|
|
|
|
def save_images(
|
|
self, images, output_path, filename_prefix, comment, extension, quality_jpeg_or_webp, lossless_webp, prompt=None, extra_pnginfo=None
|
|
) -> list[str]:
|
|
"""Process and save individual images with metadata."""
|
|
img_count = 1
|
|
paths = []
|
|
for image in images:
|
|
i = 255.0 * image.cpu().numpy()
|
|
img = Image.fromarray(np.clip(i, 0, 255).astype(np.uint8))
|
|
if images.size()[0] > 1:
|
|
filename_prefix += "_{:02d}".format(img_count)
|
|
|
|
if extension == "png":
|
|
metadata = PngInfo()
|
|
metadata.add_text("parameters", comment)
|
|
|
|
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]))
|
|
|
|
filename = f"{filename_prefix}.png"
|
|
img.save(os.path.join(output_path, filename), pnginfo=metadata, optimize=True)
|
|
else:
|
|
filename = f"{filename_prefix}.{extension}"
|
|
file = os.path.join(output_path, filename)
|
|
img.save(file, optimize=True, quality=quality_jpeg_or_webp, lossless=lossless_webp)
|
|
exif_bytes = piexif.dump(
|
|
{
|
|
"Exif": {piexif.ExifIFD.UserComment: piexif.helper.UserComment.dump(comment, encoding="unicode")},
|
|
}
|
|
)
|
|
piexif.insert(exif_bytes, file)
|
|
|
|
paths.append(filename)
|
|
img_count += 1
|
|
return paths
|
|
|
|
|
|
# Export node classes for registration
|
|
NODE_CLASSES = {
|
|
"🖥️🛠TechsToolz: Model Name Extractor": ModelNameExtractor,
|
|
"🖥️🛠TechsToolz: Save Image w/Metadata": ImageSaveWithMetadata,
|
|
}
|
|
|
|
NODE_DISPLAY_NAMES = {
|
|
"🖥️🛠TechsToolz: Model Name Extractor": "🖥️🛠TechsToolz: Model Name Extractor",
|
|
"🖥️🛠TechsToolz: Save Image w/Metadata": "🖥️🛠TechsToolz: Save Image w/Metadata",
|
|
}
|