214 lines
9.5 KiB
Python
214 lines
9.5 KiB
Python
|
"""ComfyUI -> Obsidian: сохраняет изображение и позитивный промт в заметку Obsidian.
|
|||
|
|
|
|||
|
|
Нода SaveImageObsidian:
|
|||
|
|
- сохраняет PNG в output ComfyUI (стандартная нумерация),
|
|||
|
|
- копирует PNG в папку Comfy-Promt внутри vault Obsidian,
|
|||
|
|
- создаёт .md с тем же базовым именем (без расширения):
|
|||
|
|
встроенное изображение + позитивный промт в блоке кода.
|
|||
|
|
|
|||
|
|
Режимы получения текста промта (prompt_source):
|
|||
|
|
- "string" : текст берётся со строкового входа positive_prompt.
|
|||
|
|
- "conditioning" : текст НЕ извлекается из тензора (это невозможно), а
|
|||
|
|
прослеживается по графу workflow от подключённого
|
|||
|
|
CONDITIONING-входа до породившей его ноды CLIP Text Encode
|
|||
|
|
и читается её текстовое поле.
|
|||
|
|
"""
|
|||
|
|
|
|||
|
|
import os
|
|||
|
|
import shutil
|
|||
|
|
import json
|
|||
|
|
|
|||
|
|
import numpy as np
|
|||
|
|
from PIL import Image
|
|||
|
|
from PIL.PngImagePlugin import PngInfo
|
|||
|
|
|
|||
|
|
import folder_paths
|
|||
|
|
|
|||
|
|
# =============================================================================
|
|||
|
|
# НАСТРОЙКА ПУТЕЙ / PATH CONFIGURATION
|
|||
|
|
# -----------------------------------------------------------------------------
|
|||
|
|
# RU: Чтобы изменить, куда сохраняются заметка и копия изображения, отредактируйте
|
|||
|
|
# две строки ниже.
|
|||
|
|
# VAULT_DIR — полный путь к хранилищу Obsidian (или к ЛЮБОЙ папке).
|
|||
|
|
# Используйте префикс r"" и обратные слэши на Windows.
|
|||
|
|
# OBSIDIAN_SUBFOLDER — имя подпапки внутри VAULT_DIR, куда кладутся .md и .png.
|
|||
|
|
# Если нужно писать прямо в VAULT_DIR — поставьте "".
|
|||
|
|
# После изменения сохраните файл и перезапустите ComfyUI.
|
|||
|
|
#
|
|||
|
|
# EN: To change where the note and the image copy are saved, edit the two lines
|
|||
|
|
# below.
|
|||
|
|
# VAULT_DIR — full path to the Obsidian vault (or ANY folder).
|
|||
|
|
# Use an r"" prefix and backslashes on Windows.
|
|||
|
|
# OBSIDIAN_SUBFOLDER — subfolder name inside VAULT_DIR for the .md and .png.
|
|||
|
|
# Set it to "" to write directly into VAULT_DIR.
|
|||
|
|
# Save the file and restart ComfyUI after editing.
|
|||
|
|
#
|
|||
|
|
# Примеры / Examples:
|
|||
|
|
# VAULT_DIR = r"D:\MyVault"
|
|||
|
|
# VAULT_DIR = r"\\NAS\share\Obsidian" # сетевая папка / network share
|
|||
|
|
# VAULT_DIR = "/home/user/Obsidian" # Linux / macOS
|
|||
|
|
# =============================================================================
|
|||
|
|
VAULT_DIR = r"C:\Users\dimir\Documents\ObsidianTask"
|
|||
|
|
OBSIDIAN_SUBFOLDER = "Comfy-Promt"
|
|||
|
|
# =============================================================================
|
|||
|
|
|
|||
|
|
# Имена текстовых полей, которые встречаются у разных текст-нод.
|
|||
|
|
_TEXT_KEYS = ("text", "string", "value", "prompt", "Text",
|
|||
|
|
"wildcard_text", "populated_text", "text_g", "text_l")
|
|||
|
|
|
|||
|
|
|
|||
|
|
def _resolve_string(prompt, val):
|
|||
|
|
"""val — это либо готовая строка, либо ссылка [node_id, out_idx] на ноду,
|
|||
|
|
из которой нужно достать строковое значение."""
|
|||
|
|
if isinstance(val, str):
|
|||
|
|
return val
|
|||
|
|
if isinstance(val, list) and len(val) == 2:
|
|||
|
|
src = prompt.get(str(val[0]), {})
|
|||
|
|
sinputs = src.get("inputs", {})
|
|||
|
|
for key in _TEXT_KEYS:
|
|||
|
|
if isinstance(sinputs.get(key), str):
|
|||
|
|
return sinputs[key]
|
|||
|
|
for v in sinputs.values(): # запасной вариант: любое строковое поле
|
|||
|
|
if isinstance(v, str):
|
|||
|
|
return v
|
|||
|
|
return ""
|
|||
|
|
|
|||
|
|
|
|||
|
|
def _walk_for_text(prompt, node_id, depth=0, seen=None):
|
|||
|
|
"""Рекурсивно идём по графу от node_id, собирая текст из всех найденных
|
|||
|
|
CLIP Text Encode. Промежуточные conditioning-ноды (Combine, Concat, и т.п.)
|
|||
|
|
проходятся насквозь по их входам-ссылкам."""
|
|||
|
|
if seen is None:
|
|||
|
|
seen = set()
|
|||
|
|
node_id = str(node_id)
|
|||
|
|
if depth > 30 or node_id in seen or node_id not in prompt:
|
|||
|
|
return []
|
|||
|
|
seen.add(node_id)
|
|||
|
|
|
|||
|
|
node = prompt[node_id]
|
|||
|
|
ctype = node.get("class_type", "")
|
|||
|
|
inputs = node.get("inputs", {})
|
|||
|
|
|
|||
|
|
if "CLIPTextEncode" in ctype or "CLIP Text Encode" in ctype:
|
|||
|
|
t = _resolve_string(prompt, inputs.get("text"))
|
|||
|
|
return [t] if t else []
|
|||
|
|
|
|||
|
|
texts = []
|
|||
|
|
for v in inputs.values():
|
|||
|
|
if isinstance(v, list) and len(v) == 2 and isinstance(v[0], (str, int)):
|
|||
|
|
texts.extend(_walk_for_text(prompt, v[0], depth + 1, seen))
|
|||
|
|
return texts
|
|||
|
|
|
|||
|
|
|
|||
|
|
def _text_from_conditioning(prompt, unique_id, input_name):
|
|||
|
|
"""Находит текст для CONDITIONING-входа input_name текущей ноды (unique_id)."""
|
|||
|
|
if not prompt or unique_id is None:
|
|||
|
|
return ""
|
|||
|
|
me = prompt.get(str(unique_id), {})
|
|||
|
|
link = me.get("inputs", {}).get(input_name)
|
|||
|
|
if not (isinstance(link, list) and len(link) == 2):
|
|||
|
|
return ""
|
|||
|
|
texts = _walk_for_text(prompt, link[0])
|
|||
|
|
# уникализируем, сохраняя порядок
|
|||
|
|
seen, out = set(), []
|
|||
|
|
for t in texts:
|
|||
|
|
if t and t not in seen:
|
|||
|
|
seen.add(t)
|
|||
|
|
out.append(t)
|
|||
|
|
return "\n\n".join(out)
|
|||
|
|
|
|||
|
|
|
|||
|
|
class SaveImageObsidian:
|
|||
|
|
def __init__(self):
|
|||
|
|
self.output_dir = folder_paths.get_output_directory()
|
|||
|
|
self.type = "output"
|
|||
|
|
self.compress_level = 4
|
|||
|
|
|
|||
|
|
@classmethod
|
|||
|
|
def INPUT_TYPES(cls):
|
|||
|
|
return {
|
|||
|
|
"required": {
|
|||
|
|
"images": ("IMAGE",),
|
|||
|
|
"prompt_source": (["string", "conditioning"], {"default": "string"}),
|
|||
|
|
},
|
|||
|
|
"optional": {
|
|||
|
|
"positive_prompt": ("STRING", {"forceInput": True}),
|
|||
|
|
"positive_conditioning": ("CONDITIONING",),
|
|||
|
|
"filename_prefix": ("STRING", {"default": "ComfyUI"}),
|
|||
|
|
},
|
|||
|
|
"hidden": {
|
|||
|
|
"prompt": "PROMPT",
|
|||
|
|
"extra_pnginfo": "EXTRA_PNGINFO",
|
|||
|
|
"unique_id": "UNIQUE_ID",
|
|||
|
|
},
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
RETURN_TYPES = ()
|
|||
|
|
FUNCTION = "save"
|
|||
|
|
OUTPUT_NODE = True
|
|||
|
|
CATEGORY = "Obsidian"
|
|||
|
|
|
|||
|
|
def save(self, images, prompt_source="string", positive_prompt=None,
|
|||
|
|
positive_conditioning=None, filename_prefix="ComfyUI",
|
|||
|
|
prompt=None, extra_pnginfo=None, unique_id=None):
|
|||
|
|
|
|||
|
|
# --- Определяем текст промта по выбранному режиму -------------------
|
|||
|
|
if prompt_source == "conditioning":
|
|||
|
|
prompt_text = _text_from_conditioning(prompt, unique_id,
|
|||
|
|
"positive_conditioning")
|
|||
|
|
if not prompt_text:
|
|||
|
|
prompt_text = ("[промт не найден: подключите CONDITIONING-выход "
|
|||
|
|
"ноды CLIP Text Encode напрямую]")
|
|||
|
|
else:
|
|||
|
|
prompt_text = positive_prompt if positive_prompt is not None else \
|
|||
|
|
"[промт не найден: подключите строковый вход positive_prompt]"
|
|||
|
|
|
|||
|
|
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])
|
|||
|
|
|
|||
|
|
obsidian_dir = os.path.join(VAULT_DIR, OBSIDIAN_SUBFOLDER)
|
|||
|
|
os.makedirs(obsidian_dir, exist_ok=True)
|
|||
|
|
|
|||
|
|
results = []
|
|||
|
|
for batch_number, image in enumerate(images):
|
|||
|
|
arr = 255.0 * image.cpu().numpy()
|
|||
|
|
img = Image.fromarray(np.clip(arr, 0, 255).astype(np.uint8))
|
|||
|
|
|
|||
|
|
# Метаданные workflow в PNG (как у штатной SaveImage)
|
|||
|
|
metadata = PngInfo()
|
|||
|
|
if prompt is not None:
|
|||
|
|
metadata.add_text("prompt", json.dumps(prompt))
|
|||
|
|
if extra_pnginfo is not None:
|
|||
|
|
for k in extra_pnginfo:
|
|||
|
|
metadata.add_text(k, json.dumps(extra_pnginfo[k]))
|
|||
|
|
|
|||
|
|
file = f"{filename}_{counter:05}_.png"
|
|||
|
|
out_path = os.path.join(full_output_folder, file)
|
|||
|
|
img.save(out_path, pnginfo=metadata, compress_level=self.compress_level)
|
|||
|
|
|
|||
|
|
# Копия PNG в Comfy-Promt
|
|||
|
|
shutil.copy2(out_path, os.path.join(obsidian_dir, file))
|
|||
|
|
|
|||
|
|
# Заметка с тем же базовым именем (без расширения)
|
|||
|
|
note_name = os.path.splitext(file)[0]
|
|||
|
|
note_body = (
|
|||
|
|
f"![[{file}]]\n\n"
|
|||
|
|
f"```\n{prompt_text}\n```\n"
|
|||
|
|
)
|
|||
|
|
with open(os.path.join(obsidian_dir, note_name + ".md"),
|
|||
|
|
"w", encoding="utf-8") as f:
|
|||
|
|
f.write(note_body)
|
|||
|
|
|
|||
|
|
results.append({"filename": file, "subfolder": subfolder, "type": self.type})
|
|||
|
|
counter += 1
|
|||
|
|
|
|||
|
|
return {"ui": {"images": results}}
|
|||
|
|
|
|||
|
|
|
|||
|
|
NODE_CLASS_MAPPINGS = {"SaveImageObsidian": SaveImageObsidian}
|
|||
|
|
NODE_DISPLAY_NAME_MAPPINGS = {"SaveImageObsidian": "Save Image + Obsidian Note"}
|
|||
|
|
|
|||
|
|
__all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS"]
|