import os import numpy as np import json from typing import Dict, Optional, TypedDict, Union, List from PIL import Image from PIL.PngImagePlugin import PngInfo from datetime import datetime import folder_paths from .modules.util import util from .modules.eagle_api import EagleAPI from .modules.prompt_info_extractor import PromptInfoExtractor FORCE_WRITE_PROMPT = False class ParamsInfo(TypedDict): format: str lossless_webp: bool save_tags: bool filename_template: str eagle_folder: str compression: int positive: str negative: str memo_text: str prompt: Optional[Dict] extra_pnginfo: Optional[Dict] class D2_SendEagle: def __init__(self): self.output_dir = folder_paths.get_output_directory() self.type = "output" self.output_folder = "" self.subfolder_name = "" self.eagle_api = None @classmethod def INPUT_TYPES(s): return { "required": { "images": ("IMAGE",), "format": (["webp", "png"],), # webpの時に可逆（lossless）不可逆（lossy）どちらにするか "lossless_webp": ( "BOOLEAN", {"default": True, "label_on": "lossless", "label_off": "lossy"}, ), # lossy の時の圧縮率 "compression": ( "INT", {"default": 80, "min": 1, "max": 100, "step": 1}, ), # プロンプトをEagleタグに保存するか "save_tags": ( "BOOLEAN", {"default": False, "label_on": "save", "label_off": "none"}, ), # 保存するファイル名 "filename_template": ( "STRING", {"multiline": False, "default":"{model}-{width}-{height}-{seed}"}, ), # Eagleフォルダ "eagle_folder": ( "STRING", {"default": ""} ), # ポジティブプロンプト "positive": ( "STRING", {"forceInput": True, "multiline": True}, ), # ネガティブプロンプト "negative": ( "STRING", {"forceInput": True, "multiline": True}, ), # その他メモ "memo_text": ( "STRING", {"multiline": True}, ), }, "hidden": {"prompt": "PROMPT", "extra_pnginfo": "EXTRA_PNGINFO"}, } RETURN_TYPES = () FUNCTION = "add_item" OUTPUT_NODE = True CATEGORY = "D2" # ###################### # ノード実行でEagleに画像を送る def add_item( self, images, format = "webp", lossless_webp = False, save_tags = True, filename_template = "{model}-{width}-{height}-{seed}", eagle_folder = "", compression = 80, positive = "", negative = "", memo_text = "", prompt: Optional[Dict] = None, extra_pnginfo: Optional[Dict] = None, ): self.output_folder, self.subfolder_name = self.get_output_folder() self.eagle_api = EagleAPI() results = list() params: ParamsInfo = { "format": format, "lossless_webp": lossless_webp, "save_tags": save_tags, "filename_template": filename_template, "eagle_folder": eagle_folder, "compression": compression, "positive": positive, "negative": negative, "memo_text": memo_text, "prompt": prompt, "extra_pnginfo": extra_pnginfo, } for image in images: results.append(self.create_image_object(image, params)) return {"ui": {"images": results}} # ###################### # イメージオブジェクトを作成 def create_image_object(self, image, params:ParamsInfo) -> dict: normalized_pixels = 255.0 * image.cpu().numpy() img = Image.fromarray(np.clip(normalized_pixels, 0, 255).astype(np.uint8)) # 画像をローカルに保存 file_name, file_full_path = self.save_image(img, params) # Eagleフォルダが指定されているならフォルダIDを取得 folder_id = self.eagle_api.find_or_create_folder(params["eagle_folder"]) # Send image to Eagle item = {"path": file_full_path, "name": file_name} item["annotation"] = util.make_annotation_text( params["positive"], params["negative"], params["memo_text"] ) if(params["save_tags"]): item["tags"] = util.get_prompt_tags(params["positive"]) _ret = self.eagle_api.add_item_from_path(data=item, folder_id=folder_id) return { "filename": file_name, "subfolder": self.subfolder_name, "type": self.type } # ###################### # 画像をローカルに保存 def save_image(self, img, params): file_name = "" file_full_path = "" filename_params = self.get_filename_params(img, params) if params["format"] == "webp": # Save webp image file file_name = self.get_filename(params["filename_template"], 'webp', filename_params) file_full_path = os.path.join(self.output_folder, file_name) exif_data = util.get_exif_from_prompt( img.getexif(), params["prompt"], params["extra_pnginfo"] ) img.save( file_full_path, quality = params["compression"], exif = exif_data, lossless = params["lossless_webp"], ) else: # Save png image file file_name = self.get_filename(params["filename_template"], 'png', filename_params) file_full_path = os.path.join(self.output_folder, file_name) metadata = PngInfo() if params["prompt"] is not None: metadata.add_text("prompt", json.dumps(params["prompt"])) if params["extra_pnginfo"] is not None: for x in params["extra_pnginfo"]: metadata.add_text(x, json.dumps(params["extra_pnginfo"][x])) img.save(file_full_path, pnginfo=metadata, compress_level=4) return file_name, file_full_path # ###################### # 画像保存パスを取得 def get_output_folder(self): subfolder_name = datetime.now().strftime("%Y-%m-%d") # 画像保存用フォルダが無ければ作成 output_folder = os.path.join(self.output_dir, subfolder_name) if not os.path.exists(output_folder): os.makedirs(output_folder) return output_folder, subfolder_name # ###################### # ファイルネーム用パラメーターを取得 def get_filename_params(self, img, params) -> dict: width, height = img.size gen_data = PromptInfoExtractor(params["prompt"]) model = "" steps = "" seed = "" if gen_data and hasattr(gen_data, 'info') and gen_data.info: model = os.path.splitext(gen_data.info.get("model_name", ""))[0] steps = gen_data.info.get("steps", "") seed = gen_data.info.get("seed", "") return { "width": width, "height": height, "model": model, "steps": steps, "seed": seed, } # ###################### # ファイルネームを取得 def get_filename(self, template:str, ext:str, filename_params) -> str: base = template.format( width = filename_params["width"], height = filename_params["height"], model = filename_params["model"], steps = filename_params["steps"], seed = filename_params["seed"], ) return f"{util.get_datetime_str_msec()}-{base}.{ext}"