From d37726b510be1aac1fef8945d13111b4982e4abf Mon Sep 17 00:00:00 2001 From: HappyOnigiri <253838257+NodeMeld@users.noreply.github.com> Date: Sat, 10 Jan 2026 08:55:24 +0900 Subject: [PATCH] feat: add meld nexus image manager backend --- Makefile | 2 +- py/image_manager/__init__.py | 12 +++ py/image_manager/api.py | 89 +++++++++++++++++ py/image_manager/database.py | 86 ++++++++++++++++ py/image_manager/nodes.py | 186 +++++++++++++++++++++++++++++++++++ tests/test_image_manager.py | 27 ++--- 6 files changed, 388 insertions(+), 14 deletions(-) create mode 100644 py/image_manager/__init__.py create mode 100644 py/image_manager/api.py create mode 100644 py/image_manager/database.py create mode 100644 py/image_manager/nodes.py diff --git a/Makefile b/Makefile index f9e7ddd..2aa903c 100644 --- a/Makefile +++ b/Makefile @@ -21,7 +21,7 @@ lint: check-only-ascii check-only-ascii: @echo "Checking for non-ASCII characters..." - @git ls-files | grep -v "README_ja.md" | grep -v "Makefile" | grep -v "py/image_manager.py" | xargs grep -nP "[^\x00-\x7f]" && (echo "Error: Non-ASCII characters found!" && exit 1) || echo "All files are ASCII (English) only." + @git ls-files | grep -v "README_ja.md" | grep -v "Makefile" | grep -v "py/image_manager/" | xargs grep -nP "[^\x00-\x7f]" && (echo "Error: Non-ASCII characters found!" && exit 1) || echo "All files are ASCII (English) only." repomix: repomix-full repomix-src repomix-tests diff --git a/py/image_manager/__init__.py b/py/image_manager/__init__.py new file mode 100644 index 0000000..19e4a16 --- /dev/null +++ b/py/image_manager/__init__.py @@ -0,0 +1,12 @@ +from . import api +from .database import init_db +from .nodes import MeldNexus + +# Initialize database on module load +init_db() + +NODE_CLASS_MAPPINGS = { + "MeldNexus": MeldNexus +} + +__all__ = ["MeldNexus", "NODE_CLASS_MAPPINGS", "api"] diff --git a/py/image_manager/api.py b/py/image_manager/api.py new file mode 100644 index 0000000..2791439 --- /dev/null +++ b/py/image_manager/api.py @@ -0,0 +1,89 @@ +import server +from aiohttp import web + +from .database import get_db_connection + + +@server.PromptServer.instance.routes.get("/meld-nexus/list") +async def list_images(request): + conn = get_db_connection() + cursor = conn.cursor() + + # Fetch images with basic info + cursor.execute( + "SELECT id, filename, subfolder, created_at FROM images WHERE is_deleted = 0 ORDER BY created_at DESC" + ) + images = cursor.fetchall() + + result_list = [] + + for img in images: + img_id, filename, subfolder, created_at = img + + # Fetch positive prompt + cursor.execute(""" + SELECT pp.name FROM positive_prompts pp + JOIN positive_prompt_image_relations r ON pp.id = r.positive_prompt_id + WHERE r.image_id = ? + """, (img_id,)) + pos_rows = cursor.fetchall() + positive = ", ".join([r[0] for r in pos_rows]) + + # Fetch negative prompt + cursor.execute(""" + SELECT np.name FROM negative_prompts np + JOIN negative_prompt_image_relations r ON np.id = r.negative_prompt_id + WHERE r.image_id = ? + """, (img_id,)) + neg_rows = cursor.fetchall() + negative = ", ".join([r[0] for r in neg_rows]) + + # Fetch tags + cursor.execute(""" + SELECT t.name FROM tags t + JOIN tag_image_relations r ON t.id = r.tag_id + WHERE r.image_id = ? + """, (img_id,)) + tag_rows = cursor.fetchall() + tags = [t[0] for t in tag_rows] + + result_list.append({ + "id": img_id, + "filename": filename, + "subfolder": subfolder, + "created_at": created_at, + "positive": positive, + "negative": negative, + "tags": tags + }) + + conn.close() + return web.json_response(result_list) + +@server.PromptServer.instance.routes.post("/meld-nexus/delete") +async def delete_image(request): + try: + data = await request.json() + image_id = data.get("id") + filename = data.get("filename") # Fallback for old frontend + + if not image_id and not filename: + return web.json_response({"error": "id or filename is required"}, status=400) + + conn = get_db_connection() + cursor = conn.cursor() + + if image_id: + cursor.execute("UPDATE images SET is_deleted = 1 WHERE id = ?", (image_id,)) + else: + cursor.execute("UPDATE images SET is_deleted = 1 WHERE filename = ?", (filename,)) + + if cursor.rowcount == 0: + conn.close() + return web.json_response({"error": "Image not found"}, status=404) + + conn.commit() + conn.close() + return web.json_response({"success": True}) + except Exception as e: + return web.json_response({"error": str(e)}, status=500) diff --git a/py/image_manager/database.py b/py/image_manager/database.py new file mode 100644 index 0000000..32e4624 --- /dev/null +++ b/py/image_manager/database.py @@ -0,0 +1,86 @@ +import os +import sqlite3 + +# --- DB Settings --- +DATA_DIR = os.path.join(os.path.dirname(os.path.dirname(os.path.dirname(__file__))), "data") +os.makedirs(DATA_DIR, exist_ok=True) +DB_PATH = os.path.join(DATA_DIR, "default.db") + +def init_db(): + conn = sqlite3.connect(DB_PATH) + cursor = conn.cursor() + cursor.execute(''' + CREATE TABLE IF NOT EXISTS images ( + id INTEGER PRIMARY KEY AUTOINCREMENT, + filename TEXT, + subfolder TEXT, + created_at REAL + ) + ''') + + # Add is_deleted column if not exists + try: + cursor.execute("ALTER TABLE images ADD COLUMN is_deleted INTEGER DEFAULT 0") + except sqlite3.OperationalError: + pass + + # Normalized Tables + cursor.execute(''' + CREATE TABLE IF NOT EXISTS positive_prompts ( + id INTEGER PRIMARY KEY AUTOINCREMENT, + name TEXT UNIQUE + ) + ''') + cursor.execute(''' + CREATE TABLE IF NOT EXISTS negative_prompts ( + id INTEGER PRIMARY KEY AUTOINCREMENT, + name TEXT UNIQUE + ) + ''') + cursor.execute(''' + CREATE TABLE IF NOT EXISTS tags ( + id INTEGER PRIMARY KEY AUTOINCREMENT, + name TEXT UNIQUE + ) + ''') + + # Relations + cursor.execute(''' + CREATE TABLE IF NOT EXISTS positive_prompt_image_relations ( + id INTEGER PRIMARY KEY AUTOINCREMENT, + image_id INTEGER, + positive_prompt_id INTEGER, + FOREIGN KEY(image_id) REFERENCES images(id), + FOREIGN KEY(positive_prompt_id) REFERENCES positive_prompts(id) + ) + ''') + cursor.execute(''' + CREATE TABLE IF NOT EXISTS negative_prompt_image_relations ( + id INTEGER PRIMARY KEY AUTOINCREMENT, + image_id INTEGER, + negative_prompt_id INTEGER, + FOREIGN KEY(image_id) REFERENCES images(id), + FOREIGN KEY(negative_prompt_id) REFERENCES negative_prompts(id) + ) + ''') + cursor.execute(''' + CREATE TABLE IF NOT EXISTS tag_image_relations ( + id INTEGER PRIMARY KEY AUTOINCREMENT, + image_id INTEGER, + tag_id INTEGER, + FOREIGN KEY(image_id) REFERENCES images(id), + FOREIGN KEY(tag_id) REFERENCES tags(id) + ) + ''') + + # Indices + cursor.execute("CREATE INDEX IF NOT EXISTS idx_images_created_at ON images(created_at)") + cursor.execute("CREATE INDEX IF NOT EXISTS idx_pp_name ON positive_prompts(name)") + cursor.execute("CREATE INDEX IF NOT EXISTS idx_np_name ON negative_prompts(name)") + cursor.execute("CREATE INDEX IF NOT EXISTS idx_tag_name ON tags(name)") + + conn.commit() + conn.close() + +def get_db_connection(): + return sqlite3.connect(DB_PATH) diff --git a/py/image_manager/nodes.py b/py/image_manager/nodes.py new file mode 100644 index 0000000..d50a1b5 --- /dev/null +++ b/py/image_manager/nodes.py @@ -0,0 +1,186 @@ +import json +import os +import re +import time +from datetime import datetime + +import folder_paths +import numpy as np +import server +from comfy.cli_args import args +from PIL import Image +from PIL.PngImagePlugin import PngInfo + +from ..load_image_configs.metadata_helper import MetadataHelper +from .database import get_db_connection + + +# --- Custom Node Definition --- +class MeldNexus: + def __init__(self): + self.output_dir = folder_paths.get_output_directory() + self.type = "output" + self.prefix_append = "" + self.compress_level = 4 + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "images": ("IMAGE", ), + "filename_prefix": ("STRING", { + "default": "MeldFlow", + "tooltip": ( + "The prefix for the file to save. This may include formatting " + "information such as %date:yyyy-MM-dd% or %date:yyyy_MM_dd_HHmmss%.\n" + "yyyy: year, MM: month, dd: day, HH: hour, mm: minute, ss: second." + ), + }), + }, + "optional": { + "positive": ("STRING", {"forceInput": True, "multiline": True}), + "negative": ("STRING", {"forceInput": True, "multiline": True}), + "tags": ("STRING", {"multiline": False}), + }, + "hidden": { + "prompt": "PROMPT", "extra_pnginfo": "EXTRA_PNGINFO" + }, + } + + RETURN_TYPES = () + FUNCTION = "save_images" + OUTPUT_NODE = True + CATEGORY = "MeldFlow/Nexus" + + def save_images( + self, + images, + filename_prefix="MeldFlow", + positive=None, + negative=None, + tags=None, + prompt=None, + extra_pnginfo=None, + ): + if filename_prefix is None: + filename_prefix = "MeldFlow" + + filename_prefix += self.prefix_append + + tokens = re.findall(r'%(.*?)%', filename_prefix) + for token in tokens: + if token.startswith('date:'): + format_str = token[5:] + # Simple mapper for common date formats + py_format = format_str.replace('yyyy', '%Y').replace('yy', '%y') \ + .replace('MM', '%m').replace('dd', '%d') \ + .replace('HH', '%H').replace('mm', '%M') \ + .replace('ss', '%S') + try: + formatted_date = datetime.now().strftime(py_format) + filename_prefix = filename_prefix.replace(f'%{token}%', formatted_date) + except Exception: + pass + elif token == 'date': + formatted_date = datetime.now().strftime('%Y-%m-%d') + filename_prefix = filename_prefix.replace('%date%', formatted_date) + + 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] + ) + results = list() + conn = get_db_connection() + cursor = conn.cursor() + + # Resolve Metadata + resolved_positive = positive + resolved_negative = negative + + if not resolved_positive or not resolved_negative: + # Try to parse from workflow + logs: list[str] = [] + wf_pos, wf_neg = "", "" + + if extra_pnginfo and 'workflow' in extra_pnginfo: + wf_pos, wf_neg = MetadataHelper.parse_workflow_json(extra_pnginfo['workflow'], logs) + elif prompt: + wf_pos, wf_neg = MetadataHelper.parse_prompt_json(prompt, logs) + + if not resolved_positive: + resolved_positive = wf_pos + if not resolved_negative: + resolved_negative = wf_neg + + tag_list = [] + if tags: + tag_list = [t.strip() for t in tags.split(',') if t.strip()] + + # Split prompts into lists for separate registration + pos_list = [p.strip() for p in resolved_positive.split(',') if p.strip()] if resolved_positive else [] + neg_list = [n.strip() for n in resolved_negative.split(',') if n.strip()] if resolved_negative else [] + + for (batch_number, image) in enumerate(images): + # Tensor [B, H, W, C] -> PIL + # ComfyUI images are [B, H, W, C] Tensors + i = 255. * image.cpu().numpy() + img = Image.fromarray(np.clip(i, 0, 255).astype(np.uint8)) + + # Save metadata to PNG + metadata = None + if not args.disable_metadata: + metadata = PngInfo() + if prompt is not None: + metadata.add_text("prompt", json.dumps(prompt)) + if extra_pnginfo is not None: + for x in extra_pnginfo: + metadata.add_text(x, json.dumps(extra_pnginfo[x])) + + filename_with_batch_num = filename.replace("%batch_num%", str(batch_number)) + file = f"{filename_with_batch_num}_{counter:05}.png" + img.save(os.path.join(full_output_folder, file), pnginfo=metadata, compress_level=self.compress_level) + + timestamp = time.time() + + # Insert Image + cursor.execute( + "INSERT INTO images (filename, subfolder, created_at, is_deleted) VALUES (?, ?, ?, 0)", + (file, subfolder, timestamp) + ) + image_id = cursor.lastrowid + + # Insert Prompts + for p in pos_list: + cursor.execute("INSERT OR IGNORE INTO positive_prompts (name) VALUES (?)", (p,)) + cursor.execute("SELECT id FROM positive_prompts WHERE name = ?", (p,)) + pp_id = cursor.fetchone()[0] + cursor.execute( + "INSERT INTO positive_prompt_image_relations (image_id, positive_prompt_id) VALUES (?, ?)", + (image_id, pp_id) + ) + + for n in neg_list: + cursor.execute("INSERT OR IGNORE INTO negative_prompts (name) VALUES (?)", (n,)) + cursor.execute("SELECT id FROM negative_prompts WHERE name = ?", (n,)) + np_id = cursor.fetchone()[0] + cursor.execute( + "INSERT INTO negative_prompt_image_relations (image_id, negative_prompt_id) VALUES (?, ?)", + (image_id, np_id) + ) + + # Insert Tags + for tag in tag_list: + cursor.execute("INSERT OR IGNORE INTO tags (name) VALUES (?)", (tag,)) + cursor.execute("SELECT id FROM tags WHERE name = ?", (tag,)) + tag_id = cursor.fetchone()[0] + cursor.execute("INSERT INTO tag_image_relations (image_id, tag_id) VALUES (?, ?)", (image_id, tag_id)) + + results.append({"filename": file, "subfolder": subfolder, "type": self.type}) + counter += 1 + + conn.commit() + conn.close() + + # Notify the frontend that images have been saved (Rule 2: Real-time update) + server.PromptServer.instance.send_sync("meld-nexus-image-saved", {"count": len(results)}) + + return {"ui": {"images": results}} diff --git a/tests/test_image_manager.py b/tests/test_image_manager.py index 5d68616..12f701d 100644 --- a/tests/test_image_manager.py +++ b/tests/test_image_manager.py @@ -23,7 +23,8 @@ from comfy.cli_args import args # noqa: E402 args.disable_metadata = False # Import test target -import py.image_manager # noqa: E402 +import py.image_manager.database # noqa: E402 +import py.image_manager.nodes # noqa: E402 class TestImageManager(unittest.TestCase): @@ -32,23 +33,23 @@ class TestImageManager(unittest.TestCase): self.test_dir = tempfile.mkdtemp() self.db_path = os.path.join(self.test_dir, "test.db") - # Patch DB_PATH in image_manager - self.patcher_db = patch('py.image_manager.DB_PATH', self.db_path) + # Patch DB_PATH in database module + self.patcher_db = patch('py.image_manager.database.DB_PATH', self.db_path) self.patcher_db.start() # Initialize DB - py.image_manager.init_db() + py.image_manager.database.init_db() - # Mock folder_paths + # Mock folder_paths in nodes module self.mock_output_dir = os.path.join(self.test_dir, "output") os.makedirs(self.mock_output_dir, exist_ok=True) - py.image_manager.folder_paths.get_output_directory.return_value = self.mock_output_dir + py.image_manager.nodes.folder_paths.get_output_directory.return_value = self.mock_output_dir # Reset get_save_image_path mock - py.image_manager.folder_paths.get_save_image_path.reset_mock(side_effect=True, return_value=True) + py.image_manager.nodes.folder_paths.get_save_image_path.reset_mock(side_effect=True, return_value=True) # Instance of node - self.node = py.image_manager.MeldNexus() + self.node = py.image_manager.nodes.MeldNexus() def tearDown(self): self.patcher_db.stop() @@ -57,7 +58,7 @@ class TestImageManager(unittest.TestCase): def test_save_images_filename_prefix_passed(self): """Verify that filename_prefix is correctly passed to get_save_image_path""" images = torch.zeros((1, 64, 64, 3)) - mock_get_save_path = py.image_manager.folder_paths.get_save_image_path + mock_get_save_path = py.image_manager.nodes.folder_paths.get_save_image_path mock_get_save_path.return_value = ( self.mock_output_dir, "MeldFlow", 1, "", "MeldFlow" ) @@ -72,7 +73,7 @@ class TestImageManager(unittest.TestCase): def test_save_images_db_registration(self): """Verify that image is registered in DB with correct filename""" images = torch.zeros((1, 64, 64, 3)) - mock_get_save_path = py.image_manager.folder_paths.get_save_image_path + mock_get_save_path = py.image_manager.nodes.folder_paths.get_save_image_path mock_get_save_path.return_value = ( self.mock_output_dir, "TestFile", 5, "sub", "Test" ) @@ -94,7 +95,7 @@ class TestImageManager(unittest.TestCase): def test_filename_prefix_with_batch_num_token(self): """Verify that %batch_num% in the filename returned by get_save_image_path is resolved""" images = torch.zeros((2, 64, 64, 3)) # Batch of 2 - mock_get_save_path = py.image_manager.folder_paths.get_save_image_path + mock_get_save_path = py.image_manager.nodes.folder_paths.get_save_image_path mock_get_save_path.return_value = ( self.mock_output_dir, "Batch_%batch_num%", 1, "", "Batch" ) @@ -115,7 +116,7 @@ class TestImageManager(unittest.TestCase): def test_save_images_format_fix(self): """Verify that the trailing underscore is removed from the filename""" images = torch.zeros((1, 64, 64, 3)) - mock_get_save_path = py.image_manager.folder_paths.get_save_image_path + mock_get_save_path = py.image_manager.nodes.folder_paths.get_save_image_path mock_get_save_path.return_value = ( self.mock_output_dir, "TestFile", 1, "", "Test" ) @@ -135,7 +136,7 @@ class TestImageManager(unittest.TestCase): def test_filename_prefix_date_resolution(self): """Verify that %date:format% and %date% tokens are correctly resolved""" images = torch.zeros((1, 64, 64, 3)) - mock_get_save_path = py.image_manager.folder_paths.get_save_image_path + mock_get_save_path = py.image_manager.nodes.folder_paths.get_save_image_path mock_get_save_path.return_value = (self.mock_output_dir, "resolved", 1, "", "resolved") from datetime import datetime