diff --git a/README.md b/README.md index 07087ea..6960adc 100644 --- a/README.md +++ b/README.md @@ -1,2 +1,4 @@ # Comfy-SaveCivitAI A custom node allowing to save images with CIVITAI readable datas + +_This version is based on use of_ [efficiency-nodes-comfyui](https://github.com/jags111/efficiency-nodes-comfyui) diff --git a/civitai_datas.py b/civitai_datas.py index c133f65..72b8063 100644 --- a/civitai_datas.py +++ b/civitai_datas.py @@ -29,7 +29,7 @@ class LoraData: class GenerationData: - def __init__(self, datas, loraNode, loaderNode, samplerNode): + def __init__(self, datas, loraNode, loaderNode, samplerNode, genDatas=None): self.preparePaths() self.loraStackerNodeID = None @@ -38,30 +38,15 @@ class GenerationData: self.loraNodeName = loraNode self.loaderNodeName = loaderNode self.samplerNodeName = samplerNode - - self.getNodesIds(datas) self.loras = [] - self.getPrompts(datas) - self.getSettings(datas) - self.getModel(datas) - self.getVAE(datas) - self.getLoras(datas) - - self.tiHashes = [] - - def __init__(self, datas, loraNode, loaderNode, samplerNode): - self.preparePaths() - - self.loraStackerNodeID = None - self.efficientLoaderNodeID = None - self.samplerNodeID = None - self.loraNodeName = loraNode - self.loaderNodeName = loaderNode - self.samplerNodeName = samplerNode + # If previous datas, copy datas to actual one and retrieve custom seed + if genDatas: + self.__dict__ = genDatas.__dict__.copy() + self.seed = datas.get(self.samplerNodeID, {}).get('inputs', {}).get('seed') + else: self.getNodesIds(datas) - self.loras = [] self.getPrompts(datas) self.getSettings(datas) @@ -69,7 +54,8 @@ class GenerationData: self.getVAE(datas) self.getLoras(datas) - self.tiHashes = [] + self.tiHashes = [] + def getPrompts(self, datas): self.positivePrompt = datas.get(self.efficientLoaderNodeID, {}).get('inputs', {}).get('positive') @@ -197,7 +183,6 @@ class GenerationData: if self.samplerNodeName in datas.get(node, {}).get('class_type', {}): self.samplerNodeID = node - def __str__(self): has_loras = len(self.loras) > 0 diff --git a/save_civitai_node.py b/save_civitai_node.py index 2d5f868..8872e9f 100644 --- a/save_civitai_node.py +++ b/save_civitai_node.py @@ -8,11 +8,13 @@ import json import os from datetime import datetime - from .civitai_datas import GenerationData class SaveCivitai: + + global FIRST_GEN + def __init__(self): self.output_dir = folder_paths.get_output_directory() self.type = "output" @@ -41,12 +43,26 @@ class SaveCivitai: CATEGORY = "MyTest" - def process_metadata(self, lora_stacker, efficient_loader, sampler, datas): - generated = GenerationData(datas, lora_stacker, efficient_loader, sampler) + # Retrieve metadatas from Images, to prevent model loading calculation at each batch images, + # datas are generated at first images, and copied to each next image. + # Only seeds will be retrieved, as it may be unique + def process_metadata(self, lora_stacker, efficient_loader, sampler, datas, batch_number): + global FIRST_GEN + + if batch_number == 0: + print("Generating First Data") + FIRST_GEN = GenerationData(datas, lora_stacker, efficient_loader, sampler) + generated = FIRST_GEN + + elif batch_number > 0 and FIRST_GEN: + print("Loading First Data") + generated = GenerationData(datas, lora_stacker, efficient_loader, sampler, FIRST_GEN) return generated.__str__() + # Save image with metadatas, retreived with the help of nodes names. def save_images(self, images, filename_prefix="ComfyUI", lora_stacker="LoRA Stacker", efficient_loader="Efficient Loader", sampler="KSampler", prompt=None, extra_pnginfo=None): + # Save location based on datetime, like on WEBUI filename_prefix = datetime.now().strftime("%Y-%m-%d/") + filename_prefix + self.prefix_append full_output_folder, filename, counter, subfolder, filename_prefix = folder_paths.get_save_image_path( filename_prefix, self.output_dir, images[0].shape[1], images[0].shape[0]) @@ -57,7 +73,7 @@ class SaveCivitai: metadata = PngInfo() if prompt is not None: metadata.add_text("prompt", json.dumps(prompt)) - metadata.add_text("parameters", self.process_metadata(lora_stacker, efficient_loader, sampler, prompt)) + metadata.add_text("parameters", self.process_metadata(lora_stacker, efficient_loader, sampler, prompt, batch_number)) if extra_pnginfo is not None: for x in extra_pnginfo: metadata.add_text(x, json.dumps(extra_pnginfo[x]))