nodes/HappyOnigiri-ComfyUI-Meld
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

feat: add meld nexus image manager backend

Browse Source
This commit is contained in:
HappyOnigiri
2026-01-10 08:55:24 +09:00
parent cfa21ff51f
commit d37726b510
6 changed files with 388 additions and 14 deletions
Download Patch File Download Diff File Expand all files Collapse all files
Makefile
+1 -1
View File
@@ -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
py/image_manager/__init__.py
+12
View File
@@ -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"]
py/image_manager/api.py
+89
View File
@@ -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)
py/image_manager/database.py
+86
View File
@@ -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)
py/image_manager/nodes.py
+186
View File
@@ -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}}
tests/test_image_manager.py
+14 -13
View File
@@ -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