Initial commit

2025-12-06 17:56:44 +09:00
parent 3d1906eb93
commit 9bf2df3b80
5 changed files with 247 additions and 0 deletions
@@ -0,0 +1,4 @@
+__pycache__/
+*.py[cod]
+*$py.class
+.DS_Store
@@ -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/<prompt_id>/`
+
+## 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.
@@ -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"]
@@ -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}/<prompt_id>/")
@@ -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 } }