From 2398891a6e1efdc9ef72e48ba24e2e28fc89e97c Mon Sep 17 00:00:00 2001 From: e <153011054+educator-art@users.noreply.github.com> Date: Thu, 3 Oct 2024 14:34:58 +0900 Subject: [PATCH] =?UTF-8?q?=E5=88=9D=E5=9B=9E=E3=82=B3=E3=83=9F=E3=83=83?= =?UTF-8?q?=E3=83=88?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- README.md | 53 ++++++++++++++++++++++ __init__.py | 10 +++++ nodes/nodes.py | 84 ++++++++++++++++++++++++++++++++++ nodes/nodes_advanced.py | 99 +++++++++++++++++++++++++++++++++++++++++ 4 files changed, 246 insertions(+) create mode 100644 README.md create mode 100644 __init__.py create mode 100644 nodes/nodes.py create mode 100644 nodes/nodes_advanced.py diff --git a/README.md b/README.md new file mode 100644 index 0000000..357f87f --- /dev/null +++ b/README.md @@ -0,0 +1,53 @@ +# ComfyUI-Load-DirectoryFiles + +## 概要 + +[ComfyUI](https://github.com/comfyanonymous/ComfyUI)のCustomNode + +指定したディレクトリにあるプロンプト(txt)と画像(png)を読み込みます。 +indexを指定することで、指定したファイルを出力します。 + +これは、計画した画像生成をすることを想定して作りました。 +プロンプトと画像を合わせて参照することで、プロンプトだけでは描きにくい問題を緩和します。 +例えば、ControlNetやi2iを使用することでプロンプトを補うことができるため、試行回数を減らすことに繋がります。 + +## ファイル形式について + +以下のように、同じ番号で名前を付けたテキストと画像の組み合わせを1つのフォルダに格納します。なお、ネガティブプロンプトも対応させたので、名前に"_n"を付けています + +``` +# プロンプト +001.txt +002.txt +003.txt + +# 画像 +001.png +002.png +003.png + +# ネガティブプロンプト +001_n.txt +002_n.txt +003_n.txt +``` + +## 前提 + +- ComfyUIのインストールを完了していること + +## 導入 + +``` +git clone educator-art/ComfyUI-Load-DirectoryFiles ./custom_node +``` + +## 応用 + +ComfyUIの[basic_api_example.py](https://github.com/comfyanonymous/ComfyUI/blob/master/script_examples/basic_api_example.py)を参考にしてComfyUIのAPIを活用したプログラムを作成し、指定した枚数の画像生成を自動化しました。 これにより、当初の目的を実現できました。 + +※ 応用は拡張機能の範囲外となりますので、詳しい方法は割愛します。 + +## ライセンス + + - GPL-3.0 License \ No newline at end of file diff --git a/__init__.py b/__init__.py new file mode 100644 index 0000000..8543620 --- /dev/null +++ b/__init__.py @@ -0,0 +1,10 @@ +from .nodes.nodes import * +from .nodes.nodes_advanced import * + +NODE_CLASS_MAPPINGS = { + "Load Images and Prompts from Directory": LoadImageAndPromptsFromDirectory, + "Load Images and Prompts from Directory(Advanced)": LoadImageAndPromptsFromDirectory_Advanced, +} + +print("Load Class: LoadImageAndPromptsFromDirectory") +print("Load Class: LoadImageAndPromptsFromDirectory_Advanced") diff --git a/nodes/nodes.py b/nodes/nodes.py new file mode 100644 index 0000000..a57d1a1 --- /dev/null +++ b/nodes/nodes.py @@ -0,0 +1,84 @@ +import numpy as np +import torch +from PIL import Image +import os +import glob + +class LoadImageAndPromptsFromDirectory: + + def __init__(self): + + self.max_len = 0 + self.text_files_path = [] + self.text_negative_files_path = [] + self.image_files_path = [] + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "dir": ( + "STRING", + {"multiline": False, "default": "Please input Directory Path"}, + ), + "files_index": ( + "INT", + {"default": 0, "min": 0, "max": 4096}, + ), + } + } + + RETURN_TYPES = ("STRING", "STRING", "IMAGE") + FUNCTION = "load_data" + OUTPUT_NODE = True + CATEGORY = "Load Images and Prompts from Directory" + + def check_directory(self, dir): + + if not os.path.isdir(dir): + raise ValueError(f"指定されたディレクトリが存在しません: {dir}") + + image_files = sorted(glob.glob(os.path.join(dir, "*.png"))) + text_files = sorted(glob.glob(os.path.join(dir, "*.txt"))) + text_files = [f for f in text_files if not f.endswith('_n.txt')] + text_files_negative=sorted(glob.glob(os.path.join(dir, "*_n.txt"))) + + if image_files == 0 or text_files == 0: + raise RuntimeError("テキストファイルか画像ファイルが存在しません。") + + if len(image_files) != len(text_files): + raise ValueError("テキストファイルと画像ファイルの数が一致しません。") + + self.text_files_path = text_files + self.text_negative_files_path=text_files_negative + self.image_files_path = image_files + self.max_len = len(image_files) + + + def pil_to_tensor(self, image: Image.Image) -> torch.Tensor: + return torch.from_numpy(np.array(image).astype(np.float32) / 255.0).unsqueeze(0) + + def load_data(self, dir, files_index): + if self.max_len == 0: + self.check_directory(dir) + + text_name = os.path.splitext( + os.path.basename(self.text_files_path[files_index]) + )[0] + image_name = os.path.splitext( + os.path.basename(self.image_files_path[files_index]) + )[0] + + if text_name != image_name: + raise ValueError("プロンプトと画像のファイル名が一致しません。") + + with open(self.text_files_path[files_index], "r", encoding="utf-8") as file: + prompt_string = file.read() + + with open(self.text_negative_files_path[files_index], "r", encoding="utf-8") as file: + prompt_negative_string = file.read() + + image = Image.open(self.image_files_path[files_index]).convert("RGB") + image_tensor = self.pil_to_tensor(image) + + return (prompt_string, prompt_negative_string, image_tensor) \ No newline at end of file diff --git a/nodes/nodes_advanced.py b/nodes/nodes_advanced.py new file mode 100644 index 0000000..6026522 --- /dev/null +++ b/nodes/nodes_advanced.py @@ -0,0 +1,99 @@ +""" +2024.10.02 Created +Automatic1111 WebUIのBreak文、ComfyUIへPromptをSplitして出力する対応 +""" +import numpy as np +import torch +from PIL import Image +import os +import glob + +class LoadImageAndPromptsFromDirectory_Advanced: + + def __init__(self): + + self.max_len = 0 + self.text_files_path = [] + self.text_negative_files_path = [] + self.image_files_path = [] + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "dir": ( + "STRING", + {"multiline": False, "default": "Please input Directory Path"}, + ), + "files_index": ( + "INT", + {"default": 0, "min": 0, "max": 4096}, + ), + } + } + + # Automatic 1111 Web UIのPromptのBreak文に対応（Break文：5件対応なので最大、6件のプロンプト文） + RETURN_TYPES = ("STRING", "STRING", "STRING", "STRING", "STRING", "STRING", "STRING", "IMAGE") + RETURN_NAMES = ("1st prompt", "2nd prompt", "3rd prompt","4th prompt","5th prompt" , "6th prompt", "negative prompt", "image") + FUNCTION = "load_data" + OUTPUT_NODE = True + CATEGORY = "Load Images and Prompts from Directory" + + def check_directory(self, dir): + + if not os.path.isdir(dir): + raise ValueError(f"指定されたディレクトリが存在しません: {dir}") + + image_files = sorted(glob.glob(os.path.join(dir, "*.png"))) + text_files = sorted(glob.glob(os.path.join(dir, "*.txt"))) + text_files = [f for f in text_files if not f.endswith('_n.txt')] + text_files_negative=sorted(glob.glob(os.path.join(dir, "*_n.txt"))) + + if image_files == 0 or text_files == 0: + raise RuntimeError("テキストファイルか画像ファイルが存在しません。") + + if len(image_files) != len(text_files): + raise ValueError("テキストファイルと画像ファイルの数が一致しません。") + + self.text_files_path = text_files + self.text_negative_files_path=text_files_negative + self.image_files_path = image_files + self.max_len = len(image_files) + + + def pil_to_tensor(self, image: Image.Image) -> torch.Tensor: + return torch.from_numpy(np.array(image).astype(np.float32) / 255.0).unsqueeze(0) + + def load_data(self, dir, files_index): + if self.max_len == 0: + self.check_directory(dir) + + text_name = os.path.splitext( + os.path.basename(self.text_files_path[files_index]) + )[0] + image_name = os.path.splitext( + os.path.basename(self.image_files_path[files_index]) + )[0] + + if text_name != image_name: + raise ValueError("プロンプトと画像のファイル名が一致しません。") + + with open(self.text_files_path[files_index], "r", encoding="utf-8") as file: + prompt_string = file.read() + + # 追記 + prompt_split_string = ["", "", "", "", "", ""] + prompt_string = prompt_string.replace("\n", "").strip() + segments = [segment for segment in prompt_string.split("BREAK,")] + for i in range(min(len(segments), len(prompt_split_string))): + prompt_split_string[i] = segments[i] + + with open(self.text_negative_files_path[files_index], "r", encoding="utf-8") as file: + prompt_negative_string = file.read() + + image = Image.open(self.image_files_path[files_index]).convert("RGB") + image_tensor = self.pil_to_tensor(image) + + # 修正 + return (prompt_split_string[0],prompt_split_string[1],prompt_split_string[2],prompt_split_string[3],prompt_split_string[4],prompt_split_string[5], prompt_negative_string, image_tensor) + \ No newline at end of file