From 9bf2df3b801b43fb7541b256ad05d01ae69ac7df Mon Sep 17 00:00:00 2001 From: t22m003 Date: Sat, 6 Dec 2025 17:37:12 +0900 Subject: [PATCH] Initial commit --- .gitignore | 4 ++ README.md | 31 +++++++++++ __init__.py | 16 ++++++ hooks.py | 127 ++++++++++++++++++++++++++++++++++++++++++++ save_with_prompt.py | 69 ++++++++++++++++++++++++ 5 files changed, 247 insertions(+) create mode 100644 .gitignore create mode 100644 README.md create mode 100644 __init__.py create mode 100644 hooks.py create mode 100644 save_with_prompt.py diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..0fff304 --- /dev/null +++ b/.gitignore @@ -0,0 +1,4 @@ +__pycache__/ +*.py[cod] +*$py.class +.DS_Store diff --git a/README.md b/README.md new file mode 100644 index 0000000..5915521 --- /dev/null +++ b/README.md @@ -0,0 +1,31 @@ +# ComfyUI Save Prompt and Output + +This custom node and hook for ComfyUI automatically saves all intermediate outputs (images and metadata) for every executed node in your workflow. + +## Features + +- **Automatic Hook**: No need to place special nodes. Just running your workflow captures everything. +- **Workflow Logging**: Saves the full workflow and parameters used for generation. +- **Intermediate Output**: Captures outputs from executed nodes, including images and non-image data (serialized). +- **Organized Storage**: Creates a unique folder for each execution (Queue) based on the Prompt ID. + - Path: `ComfyUI/output/hooked_outputs//` + +## Installation + +1. Clone this repository into your `ComfyUI/custom_nodes/` directory: + ```bash + cd ComfyUI/custom_nodes/ + git clone git@github.com:t22m003/ComfyUI_SavePromptAndOutput.git + ``` + +2. Restart ComfyUI. + +## Usage + +Just run your workflows as usual. +Check the `ComfyUI/output/hooked_outputs/` directory for results. + +## Configuration + +Currently, the save path is hardcoded to `ComfyUI/output/hooked_outputs`. +You can modify `hooks.py` if you need to change the location. diff --git a/__init__.py b/__init__.py new file mode 100644 index 0000000..27dad7d --- /dev/null +++ b/__init__.py @@ -0,0 +1,16 @@ + +from .save_with_prompt import SaveImageWithPrompt +from .hooks import apply_hook + +# Apply the execution hook +apply_hook() + +NODE_CLASS_MAPPINGS = { + "SaveImageWithPrompt": SaveImageWithPrompt +} + +NODE_DISPLAY_NAME_MAPPINGS = { + "SaveImageWithPrompt": "Save Image With Prompt" +} + +__all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS"] diff --git a/hooks.py b/hooks.py new file mode 100644 index 0000000..9e1e3a2 --- /dev/null +++ b/hooks.py @@ -0,0 +1,127 @@ + +import os +import json +import torch +import numpy as np +from PIL import Image +import folder_paths +import execution +import time + +# オリジナルのexecute関数を保持 +original_execute = execution.execute + +# ベースの保存先ディレクトリ +HOOK_BASE_DIR = os.path.join(folder_paths.get_output_directory(), "hooked_outputs") + +def get_save_dir(prompt_id): + # prompt_id (実行ID) ごとにフォルダを分ける + # タイムスタンプは実行開始時ではなくノード実行時になってしまうが、prompt_id自体がユニークなIDとして使える + # 分かりやすさのためにディレクトリ名に日付などは入れず prompt_id だけにするか、 + # どこかで開始時刻を取れれば良いが、単純化のため prompt_id をフォルダ名とする + save_dir = os.path.join(HOOK_BASE_DIR, prompt_id) + os.makedirs(save_dir, exist_ok=True) + return save_dir + +def save_image(images, filename_prefix, save_dir): + results = [] + # Tensor以外（リストなど）が来る場合もあるので再帰的に処理するか、Tensorのみ抽出 + if isinstance(images, list): + for img in images: + results.extend(save_image(img, filename_prefix, save_dir)) + return results + + if not isinstance(images, torch.Tensor): + return [] + + # 画像として保存可能なTensorかチェック (B, H, W, C) + if images.ndim == 4 and images.shape[-1] in [1, 3, 4]: + for batch_idx, img_tensor in enumerate(images): + try: + img_array = 255. * img_tensor.cpu().numpy() + img = Image.fromarray(np.clip(img_array, 0, 255).astype(np.uint8)) + + filename = f"{filename_prefix}_{batch_idx}.png" + file_path = os.path.join(save_dir, filename) + img.save(file_path, optimize=True) + results.append(filename) + except Exception as e: + # 稀に形状が合わない等のエラーがあっても全体を止めない + print(f"[Hook Image Save Error] {e}") + + return results + +def serialize_value(value): + if isinstance(value, (int, float, str, bool, type(None))): + return value + elif isinstance(value, torch.Tensor): + return f"Tensor shape={value.shape} dtype={value.dtype} device={value.device}" + elif isinstance(value, list): + return [serialize_value(x) for x in value] + elif isinstance(value, dict): + return {k: serialize_value(v) for k, v in value.items()} + else: + return str(type(value)) + +async def new_execute(server, dynprompt, caches, current_item, extra_data, executed, prompt_id, execution_list, pending_subgraph_results, pending_async_nodes, ui_outputs): + # オリジナルの実行 + result = await original_execute(server, dynprompt, caches, current_item, extra_data, executed, prompt_id, execution_list, pending_subgraph_results, pending_async_nodes, ui_outputs) + + # 実行成功時のみ処理 + if result[0] == execution.ExecutionResult.SUCCESS: + try: + # キャッシュから出力を取得 + cached = caches.outputs.get(current_item) + if cached is not None and cached.outputs: + timestamp = int(time.time()) + node_info = dynprompt.get_node(current_item) + class_type = node_info.get("class_type", "Unknown") + + # 保存ディレクトリの決定 + save_dir = get_save_dir(prompt_id) + + # ファイル名のプレフィックス (タイムスタンプ_ノードID_タイプ) + filename_base = f"{timestamp}_{current_item}_{class_type}" + + saved_files = [] + outputs_summary = [] + + # 全ての出力を走査して保存 + for out_idx, output_data in enumerate(cached.outputs): + try: + # save_dir を渡すように変更 + files = save_image(output_data, f"{filename_base}_out{out_idx}", save_dir) + if files: + saved_files.extend(files) + outputs_summary.append({"index": out_idx, "type": "image", "files": files}) + else: + val_str = serialize_value(output_data) + outputs_summary.append({"index": out_idx, "type": "other", "value": val_str}) + except Exception as e: + outputs_summary.append({"index": out_idx, "type": "error", "error": str(e)}) + + # メタデータ(実行ログ)の保存 + metadata = { + "node_id": current_item, + "class_type": class_type, + "prompt_id": prompt_id, + "timestamp": timestamp, + "outputs": outputs_summary, + "inputs": node_info.get("inputs"), + "workflow": dynprompt.get_original_prompt() + } + + json_path = os.path.join(save_dir, f"{filename_base}_meta.json") + with open(json_path, "w", encoding="utf-8") as f: + json.dump(metadata, f, indent=4, ensure_ascii=False) + + except Exception as e: + print(f"[Hook Error] Failed to save output for node {current_item}: {e}") + import traceback + traceback.print_exc() + + return result + +def apply_hook(): + execution.execute = new_execute + print(f"ComfyUI Execution Hook Applied: All intermediate outputs will be logged to {HOOK_BASE_DIR}//") diff --git a/save_with_prompt.py b/save_with_prompt.py new file mode 100644 index 0000000..07520d4 --- /dev/null +++ b/save_with_prompt.py @@ -0,0 +1,69 @@ + +import os +import json +import numpy as np +from PIL import Image +from datetime import datetime +import folder_paths + +class SaveImageWithPrompt: + def __init__(self): + self.output_dir = folder_paths.get_output_directory() + self.type = "output" + self.prefix_append = "" + + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "images": ("IMAGE", ), + "filename_prefix": ("STRING", {"default": "ComfyUI_Custom"}), + "save_prompt": ("BOOLEAN", {"default": True, "label_on": "enable", "label_off": "disable"}), + "save_extra_pnginfo": ("BOOLEAN", {"default": True, "label_on": "enable", "label_off": "disable"}), + }, + "hidden": { + "prompt": "PROMPT", + "extra_pnginfo": "EXTRA_PNGINFO" + }, + } + + RETURN_TYPES = () + FUNCTION = "save_images" + OUTPUT_NODE = True + CATEGORY = "Custom/Save" + + def save_images(self, images, filename_prefix="ComfyUI_Custom", save_prompt=True, save_extra_pnginfo=True, prompt=None, extra_pnginfo=None): + 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]) + results = list() + + # Prepare metadata to save in JSON + metadata = {} + if save_prompt and prompt is not None: + metadata["prompt"] = prompt + if save_extra_pnginfo and extra_pnginfo is not None: + metadata["extra_pnginfo"] = extra_pnginfo + + # Save metadata to a JSON file if there is any (and if requested) + if metadata: + metadata_filename = f"{filename}_{counter:05}_metadata.json" + metadata_path = os.path.join(full_output_folder, metadata_filename) + with open(metadata_path, "w", encoding="utf-8") as f: + json.dump(metadata, f, indent=4, ensure_ascii=False) + + for image in images: + i = 255. * image.cpu().numpy() + img = Image.fromarray(np.clip(i, 0, 255).astype(np.uint8)) + + file = f"{filename}_{counter:05}_.png" + + img.save(os.path.join(full_output_folder, file), optimize=True) + + results.append({ + "filename": file, + "subfolder": subfolder, + "type": self.type + }) + counter += 1 + + return { "ui": { "images": results } }