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"]
|