315 lines
11 KiB
Python
315 lines
11 KiB
Python
import json
|
|
import os
|
|
import numpy as np
|
|
import json
|
|
from typing import Dict, Optional, Literal
|
|
|
|
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.params_extractor import ParamsExtractor
|
|
|
|
from .my_types import TNodeParams, TGenInfo, D2_TD2Pipe
|
|
|
|
FORCE_WRITE_PROMPT = False
|
|
|
|
|
|
|
|
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:EagleAPI = EagleAPI()
|
|
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"images": ("IMAGE",),
|
|
"format": (["webp", "png", "jpeg"],),
|
|
# webpの時に可逆（lossless）不可逆（lossy）どちらにするか
|
|
"lossless_webp": (
|
|
"BOOLEAN",
|
|
{"default": True, "label_on": "lossless", "label_off": "lossy"},
|
|
),
|
|
# webp lossy または jpeg の時の圧縮率
|
|
"compression": (
|
|
"INT",
|
|
{"default": 90, "min": 1, "max": 100, "step": 1},
|
|
),
|
|
# プロンプトやモデルをEagleタグに保存するか
|
|
"save_tags": ([
|
|
"None",
|
|
"Prompt + Checkpoint",
|
|
"Prompt",
|
|
"Checkpoint",
|
|
],),
|
|
# 保存するファイル名
|
|
"filename_template": (
|
|
"STRING",
|
|
{"multiline": False, "default":"{model}-{seed}"},
|
|
),
|
|
# Eagleフォルダ
|
|
"eagle_folder": (
|
|
"STRING",
|
|
{"default": ""}
|
|
),
|
|
# プレビュー表示するか
|
|
"preview": (
|
|
"BOOLEAN",
|
|
{"default": True, "label_on": "ON", "label_off": "OFF"},
|
|
),
|
|
},
|
|
"optional":{
|
|
# ポジティブプロンプト
|
|
"positive": (
|
|
"STRING",
|
|
{"forceInput": True, "multiline": True},
|
|
),
|
|
# ネガティブプロンプト
|
|
"negative": (
|
|
"STRING",
|
|
{"forceInput": True, "multiline": True},
|
|
),
|
|
# その他メモ
|
|
"memo_text": (
|
|
"STRING",
|
|
{"multiline": True},
|
|
),
|
|
"d2_pipe": ("D2_TD2Pipe",),
|
|
},
|
|
"hidden": {"prompt": "PROMPT", "extra_pnginfo": "EXTRA_PNGINFO"},
|
|
}
|
|
|
|
RETURN_TYPES = ("STRING", "STRING", "IMAGE")
|
|
RETURN_NAMES = ("positive", "negative", "IMAGE", )
|
|
FUNCTION = "add_item"
|
|
OUTPUT_NODE = True
|
|
CATEGORY = "D2"
|
|
|
|
# ######################
|
|
# ノード実行で呼ばれる関数
|
|
# Eagleに画像を送る
|
|
def add_item(
|
|
self,
|
|
images,
|
|
format = "webp",
|
|
lossless_webp = False,
|
|
save_tags = "None",
|
|
filename_template = "{model}-{width}-{height}-{seed}",
|
|
eagle_folder = "",
|
|
compression = 80,
|
|
positive = "",
|
|
negative = "",
|
|
preview = True,
|
|
memo_text = "",
|
|
d2_pipe: Optional[D2_TD2Pipe] = None,
|
|
prompt: Optional[Dict] = None,
|
|
extra_pnginfo: Optional[Dict] = None,
|
|
):
|
|
self.output_folder, self.subfolder_name = self.get_output_folder()
|
|
|
|
results = list()
|
|
params = TNodeParams(
|
|
format = format,
|
|
lossless_webp = lossless_webp,
|
|
save_tags = save_tags,
|
|
filename_template = filename_template,
|
|
eagle_folder = eagle_folder,
|
|
compression = compression,
|
|
positive = self.__class__.get_prompt_value("positive", positive, d2_pipe),
|
|
negative = self.__class__.get_prompt_value("negative", negative, d2_pipe),
|
|
memo_text = memo_text,
|
|
prompt = prompt,
|
|
extra_pnginfo = extra_pnginfo,
|
|
)
|
|
|
|
for image in images:
|
|
results.append(self.create_image_object(image, params, d2_pipe))
|
|
|
|
if(preview):
|
|
return {
|
|
"ui": {"images": results},
|
|
"result": (params["positive"], params["negative"], images,)
|
|
}
|
|
|
|
return {
|
|
"result": (params["positive"], params["negative"], images,)
|
|
}
|
|
|
|
@classmethod
|
|
def get_prompt_value(cls, mode: Literal["positive","negative"], value: Optional[str], d2_pipe: Optional[D2_TD2Pipe]) -> str:
|
|
if value:
|
|
return value
|
|
if d2_pipe is not None and getattr(d2_pipe, mode) is not None:
|
|
return getattr(d2_pipe, mode)
|
|
return ""
|
|
|
|
# ######################
|
|
# イメージオブジェクトを作成
|
|
def create_image_object(self, image, params:TNodeParams, d2_pipe:D2_TD2Pipe | None) -> dict:
|
|
normalized_pixels = 255.0 * image.cpu().numpy()
|
|
img = Image.fromarray(np.clip(normalized_pixels, 0, 255).astype(np.uint8))
|
|
|
|
# 生成パラメータ整理
|
|
paramsExtractor = self.create_generate_params(img, params, d2_pipe)
|
|
# 必要な生成パラメーターをまとめたもの
|
|
gen_info = paramsExtractor.gen_info
|
|
# EagleやPNGInfo用に整形したもの
|
|
formated_info = paramsExtractor.format_info(params["memo_text"])
|
|
|
|
# print("generate_params", gen_info)
|
|
# print("format_info", formated_info)
|
|
|
|
# 画像をローカルに保存
|
|
file_name, file_full_path = self.save_image(img, params, gen_info, formated_info)
|
|
|
|
# Eagleフォルダが指定されているならフォルダIDを取得
|
|
folder_id = self.eagle_api.find_or_create_folder(params["eagle_folder"])
|
|
|
|
# Eagleに送る情報を作成
|
|
item = {
|
|
"path": file_full_path,
|
|
"name": file_name,
|
|
"annotation": formated_info,
|
|
"tags": [],
|
|
}
|
|
|
|
# タグを取得
|
|
item["tags"] = self.get_tags(params, gen_info)
|
|
|
|
_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 get_tags(self, params:TNodeParams, gen_info:TGenInfo) -> list:
|
|
if(params["save_tags"] == "Prompt + Checkpoint"):
|
|
return [*util.get_prompt_tags(gen_info["positive"]), gen_info["model_name"]]
|
|
|
|
elif(params["save_tags"] == "Prompt"):
|
|
return util.get_prompt_tags(gen_info["positive"])
|
|
|
|
elif(params["save_tags"] == "Checkpoint"):
|
|
return [gen_info["model_name"]]
|
|
|
|
return []
|
|
|
|
|
|
# ######################
|
|
# 画像をローカルに保存
|
|
def save_image(self, img, params:TNodeParams, gen_info:TGenInfo, formated_info:str):
|
|
file_name = ""
|
|
file_full_path = ""
|
|
|
|
if params["format"] == "webp":
|
|
# Save webp image file
|
|
file_name = self.get_filename(params["filename_template"], 'webp', gen_info)
|
|
file_full_path = os.path.join(self.output_folder, file_name)
|
|
|
|
exif = util.get_exif_from_prompt(img, formated_info, params["extra_pnginfo"], params["prompt"])
|
|
|
|
img.save(
|
|
file_full_path,
|
|
quality = params["compression"],
|
|
exif = exif,
|
|
lossless = params["lossless_webp"],
|
|
)
|
|
|
|
elif params["format"] == "jpeg":
|
|
# Save jpeg image file
|
|
file_name = self.get_filename(params["filename_template"], 'jpeg', gen_info)
|
|
file_full_path = os.path.join(self.output_folder, file_name)
|
|
|
|
exif = util.get_exif_from_prompt(img, formated_info, params["extra_pnginfo"], params["prompt"])
|
|
|
|
img.save(
|
|
file_full_path,
|
|
quality = params["compression"],
|
|
optimize = True,
|
|
exif = exif,
|
|
)
|
|
|
|
else:
|
|
# Save png image file
|
|
file_name = self.get_filename(params["filename_template"], 'png', gen_info)
|
|
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 create_generate_params(self, img, params:TNodeParams, d2_pipe:D2_TD2Pipe | None) -> ParamsExtractor:
|
|
# print("[SendEagle] create_generate_params - ", params )
|
|
paramsExtractor = ParamsExtractor(params)
|
|
paramsExtractor.gen_info["width"] = img.width
|
|
paramsExtractor.gen_info["height"] = img.height
|
|
|
|
|
|
# pipe が指定されていればその値を入力する
|
|
if d2_pipe != None:
|
|
if d2_pipe.steps != None:
|
|
paramsExtractor.gen_info["steps"] = d2_pipe.steps
|
|
if d2_pipe.sampler_name:
|
|
paramsExtractor.gen_info["sampler_name"] = d2_pipe.sampler_name
|
|
if d2_pipe.scheduler:
|
|
paramsExtractor.gen_info["scheduler"] = d2_pipe.scheduler
|
|
if d2_pipe.cfg != None:
|
|
paramsExtractor.gen_info["cfg"] = d2_pipe.cfg
|
|
if d2_pipe.seed != None:
|
|
paramsExtractor.gen_info["seed"] = d2_pipe.seed
|
|
if d2_pipe.ckpt_name:
|
|
paramsExtractor.gen_info["model_name"] = d2_pipe.ckpt_name.replace("\\", "__")
|
|
|
|
return paramsExtractor
|
|
|
|
|
|
# ######################
|
|
# ファイルネームを取得
|
|
def get_filename(self, template:str, ext:str, gen_info:TGenInfo) -> str:
|
|
base = template.format(
|
|
width = gen_info["width"],
|
|
height = gen_info["height"],
|
|
model = gen_info["model_name"],
|
|
steps = gen_info["steps"],
|
|
seed = gen_info["seed"],
|
|
)
|
|
|
|
return f"{util.get_datetime_str_msec()}-{base}.{ext}"
|