From 72fd1ddc73ab7664eeaa75686c3062adb6ecf71e Mon Sep 17 00:00:00 2001 From: Raykosan <71210053+Raykosan@users.noreply.github.com> Date: Thu, 6 Aug 2026 15:05:15 +0300 Subject: [PATCH] Added the RS Decode Save node --- Rayko_VAE_Save_Image.py | 237 +++++++++++++++++++++++++++++++++++++++ __init__.py | 2 +- package.json | 3 +- pyproject.toml | 2 +- web/rs_vae_save_image.js | 202 +++++++++++++++++++++++++++++++++ 5 files changed, 443 insertions(+), 3 deletions(-) create mode 100644 Rayko_VAE_Save_Image.py create mode 100644 web/rs_vae_save_image.js diff --git a/Rayko_VAE_Save_Image.py b/Rayko_VAE_Save_Image.py new file mode 100644 index 0000000..42caf95 --- /dev/null +++ b/Rayko_VAE_Save_Image.py @@ -0,0 +1,237 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright 2025-2026 Raykosan (RaykoStudio) +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import os +import json +import numpy as np +import re +from datetime import datetime +from PIL import Image +from PIL.PngImagePlugin import PngInfo +import folder_paths +import torch + + +class RS_VAE_Decode_Save: + def __init__(self): + self.output_dir = folder_paths.get_output_directory() + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "samples": ("LATENT",), + "vae": ("VAE",), + "filename_prefix": ("STRING", {"default": "ComfyUI"}), + "format": (["png", "jpg", "webp"], {"default": "png"}), + }, + "hidden": { + "prompt": "PROMPT", + "extra_pnginfo": "EXTRA_PNGINFO", + }, + } + + RETURN_TYPES = ("IMAGE",) + RETURN_NAMES = ("image",) + FUNCTION = "decode_and_save" + OUTPUT_NODE = True + CATEGORY = "🦊 RaykoStudio" + DESCRIPTION = "Combines native VAE Decode and Save Image into a single node." + + PNG_COMPRESSION = 1 + JPG_QUALITY = 90 + WEBP_QUALITY = 90 + EMBED_WORKFLOW = True + OVERWRITE_EXISTING = False + + @staticmethod + def _normalize_images(images): + while images.dim() > 4 and images.shape[0] == 1: + images = images.squeeze(0) + if images.dim() == 2: + images = images.unsqueeze(0).unsqueeze(-1) + if images.dim() == 3: + images = images.unsqueeze(0) + if images.dim() != 4: + raise ValueError(f"[RS] Unexpected tensor dims: {images.shape}") + b, d1, d2, d3 = images.shape + second_is_ch = d1 in (1, 3, 4) and d3 not in (1, 3, 4) + last_is_ch = d3 in (1, 3, 4) and d1 not in (1, 3, 4) + if second_is_ch and not last_is_ch: + images = images.permute(0, 2, 3, 1) + elif d1 == 1 and d3 > 4: + images = images.squeeze(1).unsqueeze(-1) + if images.shape[-1] == 1: + images = images.repeat(1, 1, 1, 3) + if images.dtype != torch.float32: + images = images.float() + images = torch.clamp(images, 0.0, 1.0) + return images + + def _get_next_filename(self, directory, base_name, extension, is_batch, batch_index=None): + try: + if self.OVERWRITE_EXISTING: + if is_batch and batch_index is not None: + return os.path.join(directory, f"{base_name}_{batch_index:03d}.{extension}") + return os.path.join(directory, f"{base_name}_001.{extension}") + if is_batch and batch_index is not None: + base_with_index = f"{base_name}_{batch_index:03d}" + pattern = re.compile(rf'^{re.escape(base_with_index)}_(\d+)\.{re.escape(extension)}$') + else: + pattern = re.compile(rf'^{re.escape(base_name)}_(\d+)\.{re.escape(extension)}$') + max_num = 0 + if os.path.exists(directory): + for f in os.listdir(directory): + match = pattern.match(f) + if match: + num = int(match.group(1)) + if num > max_num: + max_num = num + next_num = max_num + 1 + if is_batch and batch_index is not None: + filename = f"{base_name}_{batch_index:03d}_{next_num:03d}.{extension}" + else: + filename = f"{base_name}_{next_num:03d}.{extension}" + return os.path.join(directory, filename) + except Exception as e: + print(f"[RS] Filename error: {e}") + ts = datetime.now().strftime("%Y%m%d_%H%M%S") + return os.path.join(directory, f"{base_name}_{ts}.{extension}") + + def _build_workflow_json(self, prompt, extra_pnginfo): + try: + if extra_pnginfo and 'workflow' in extra_pnginfo: + wf = extra_pnginfo['workflow'] + if 'prompt' not in wf: + wf['prompt'] = prompt + return wf + return {"prompt": prompt or {}, "workflow_info": "Generated by RS Decode Save"} + except Exception as e: + print(f"[RS] Workflow JSON error: {e}") + return {"prompt": prompt or {}} + + def _save_workflow_json(self, image_path, workflow_data): + try: + json_path = os.path.splitext(image_path)[0] + ".json" + with open(json_path, 'w', encoding='utf-8') as f: + json.dump(workflow_data, f, ensure_ascii=False, indent=2) + except Exception as e: + print(f"[RS] Workflow JSON save error: {e}") + + def _save_single(self, image_tensor, save_dir, base_filename, fmt, + prompt, extra_pnginfo, is_batch, batch_index=None): + try: + ext_map = {"png": "png", "jpg": "jpg", "webp": "webp"} + extension = ext_map.get(fmt, "png") + filepath = self._get_next_filename(save_dir, base_filename, extension, is_batch, batch_index) + img_array = np.clip(255.0 * image_tensor.cpu().numpy(), 0, 255).astype(np.uint8) + img = Image.fromarray(img_array) + save_kwargs = {} + if fmt == "png": + save_kwargs["compress_level"] = self.PNG_COMPRESSION + if self.EMBED_WORKFLOW and prompt: + metadata = PngInfo() + metadata.add_text("prompt", json.dumps(prompt)) + if extra_pnginfo: + for key, value in extra_pnginfo.items(): + try: + metadata.add_text(key, json.dumps(value)) + except Exception: + pass + save_kwargs["pnginfo"] = metadata + elif fmt == "jpg": + save_kwargs["quality"] = self.JPG_QUALITY + save_kwargs["optimize"] = True + save_kwargs["progressive"] = True + if self.EMBED_WORKFLOW and (prompt or extra_pnginfo): + self._save_workflow_json(filepath, self._build_workflow_json(prompt, extra_pnginfo)) + elif fmt == "webp": + save_kwargs["quality"] = self.WEBP_QUALITY + save_kwargs["method"] = 4 + if self.EMBED_WORKFLOW and (prompt or extra_pnginfo): + self._save_workflow_json(filepath, self._build_workflow_json(prompt, extra_pnginfo)) + img.save(filepath, **save_kwargs) + return filepath + except Exception as e: + print(f"[RS] Save error: {e}") + return "" + + def decode_and_save(self, samples, vae, filename_prefix, format, + prompt=None, extra_pnginfo=None): + images = vae.decode(samples["samples"]) + images = self._normalize_images(images) + + full_output_folder, fn_prefix, counter, subfolder, _ = \ + folder_paths.get_save_image_path(filename_prefix, self.output_dir, + images[0].shape[1], images[0].shape[0]) + + batch_size = images.shape[0] + is_batch = batch_size > 1 + ext_map = {"png": "png", "jpg": "jpg", "webp": "webp"} + extension = ext_map.get(format, "png") + + saved_files = [] + for i in range(batch_size): + batch_index = i + 1 if is_batch else None + + file_counter = counter + i + filename = f"{fn_prefix}_{file_counter:05}.{extension}" + filepath = os.path.join(full_output_folder, filename) + + img_array = np.clip(255.0 * images[i].cpu().numpy(), 0, 255).astype(np.uint8) + img = Image.fromarray(img_array) + + save_kwargs = {} + if format == "png": + save_kwargs["compress_level"] = self.PNG_COMPRESSION + if self.EMBED_WORKFLOW and prompt: + metadata = PngInfo() + metadata.add_text("prompt", json.dumps(prompt)) + if extra_pnginfo: + for key, value in extra_pnginfo.items(): + try: + metadata.add_text(key, json.dumps(value)) + except Exception: + pass + save_kwargs["pnginfo"] = metadata + elif format == "jpg": + save_kwargs["quality"] = self.JPG_QUALITY + save_kwargs["optimize"] = True + save_kwargs["progressive"] = True + if self.EMBED_WORKFLOW and (prompt or extra_pnginfo): + self._save_workflow_json(filepath, self._build_workflow_json(prompt, extra_pnginfo)) + elif format == "webp": + save_kwargs["quality"] = self.WEBP_QUALITY + save_kwargs["method"] = 4 + if self.EMBED_WORKFLOW and (prompt or extra_pnginfo): + self._save_workflow_json(filepath, self._build_workflow_json(prompt, extra_pnginfo)) + + img.save(filepath, **save_kwargs) + + saved_files.append({ + "filename": filename, + "subfolder": subfolder, + "type": "output" + }) + + return {"ui": {"images": saved_files}, "result": (images,)} + +NODE_CLASS_MAPPINGS = { + "RS_VAE_Decode_Save": RS_VAE_Decode_Save, +} + +NODE_DISPLAY_NAME_MAPPINGS = { + "RS_VAE_Decode_Save": "🦊 RS Decode Save", +} \ No newline at end of file diff --git a/__init__.py b/__init__.py index 057b539..97f6f90 100644 --- a/__init__.py +++ b/__init__.py @@ -29,7 +29,7 @@ IGNORED_MODULES: Set[str] = {'__init__.py'} REQUIRED_MODULES: Set[str] = set() PACKAGE_DISPLAY_NAME = "ComfyUI_RaykoStudio" -__version__ = "0.35.1" +__version__ = "0.36.0" def load_modules(): print(f"🦊 {PACKAGE_DISPLAY_NAME}: Version: {__version__}") diff --git a/package.json b/package.json index e507cc4..0bd2d15 100644 --- a/package.json +++ b/package.json @@ -1,6 +1,6 @@ { "name": "ComfyUI_RaykoStudio", - "version": "0.35.1", + "version": "0.36.0", "description": "Set of custom nodes for ComfyUI providing additional image processing capabilities", "author": "Raykosan", "license": "Apache License 2.0", @@ -39,6 +39,7 @@ "Rayko_Spline_Mask.py" "Rayko_Styles_Loader.py" "Rayko_Text_Overlay_Pro.py" + "Rayko_VAE_Save_Image.py" "settings_manager_api.py" ] } diff --git a/pyproject.toml b/pyproject.toml index acab9ae..841cf34 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,7 +1,7 @@ [project] name = "ComfyUI_RaykoStudio" description = "Set of custom nodes for ComfyUI providing additional image processing capabilities" -version = "0.35.1" +version = "0.36.0" license = { file = "LICENSE" } dependencies = ["torch", "numpy", "Pillow>=10.0.0", "freetype-py>=2.5.1", "pycairo>=1.29.0", "scipy>=1.10.0"] diff --git a/web/rs_vae_save_image.js b/web/rs_vae_save_image.js new file mode 100644 index 0000000..4e4eb9f --- /dev/null +++ b/web/rs_vae_save_image.js @@ -0,0 +1,202 @@ +import { app } from "../../scripts/app.js"; + +const NODE_TYPE = "RS_VAE_Decode_Save"; + +app.registerExtension({ + name: "RaykoStudio.VAESaveImage", + + async beforeRegisterNodeDef(nodeType, nodeData, app) { + if (nodeData.name !== NODE_TYPE) return; + + const onNodeCreated = nodeType.prototype.onNodeCreated; + nodeType.prototype.onNodeCreated = function () { + const result = onNodeCreated ? onNodeCreated.apply(this, arguments) : undefined; + const self = this; + + if (this.widgets) this.widgets.forEach(w => w.hidden = true); + + this.rs_data = { filename_prefix: "ComfyUI", format: "png" }; + + const prefixW = this.widgets?.find(w => w.name === "filename_prefix"); + const formatW = this.widgets?.find(w => w.name === "format"); + if (prefixW) this.rs_data.filename_prefix = prefixW.value || "ComfyUI"; + if (formatW) this.rs_data.format = formatW.value || "png"; + + this.rowHeight = 24; + this.padding = 10; + this.labelWidth = 70; + this.targetWidth = 320; + this.clickZones = []; + this.widgetsHeight = 0; + this.imgs = []; + this.imageIndex = 0; + this.setSize([this.targetWidth, 300]); + this.widgets_start_y = 107; + + this.syncToWidgets = function () { + if (prefixW) prefixW.value = self.rs_data.filename_prefix; + if (formatW) formatW.value = self.rs_data.format; + }; + + const onExecuted = this.onExecuted; + this.onExecuted = function (message) { + const r = onExecuted ? onExecuted.apply(this, arguments) : undefined; + if (message?.images) { + this.imgs = []; + this.imageIndex = 0; + for (const image of message.images) { + const img = new Image(); + img.onload = () => { + if (this.graph) this.graph.setDirtyCanvas(true, true); + }; + img.src = `/view?filename=${encodeURIComponent(image.filename)}&type=${image.type}&subfolder=${encodeURIComponent(image.subfolder || '')}`; + this.imgs.push(img); + } + } + return r; + }; + + const origODF = this.onDrawForeground; + this.onDrawForeground = function (ctx, vr) { + if (origODF) origODF.apply(this, arguments); + this.clickZones = []; + const p = this.padding, lW = this.labelWidth, rH = this.rowHeight; + const iW = this.size[0] - p * 2 - lW; + let y = 45; + + this.drawLabel(ctx, "PREFIX", p, y, lW, rH); + this.drawStringField(ctx, this.rs_data.filename_prefix, p + lW, y, iW, rH); + this.clickZones.push({ type: "prefix", x: p + lW, y, w: iW, h: rH }); + y += rH + 4; + + this.drawLabel(ctx, "FORMAT", p, y, lW, rH); + this.drawComboField(ctx, this.rs_data.format.toUpperCase(), p + lW, y, iW, rH); + this.clickZones.push({ type: "format", x: p + lW, y, w: iW, h: rH }); + y += rH + 10; + + this.widgetsHeight = y; + }; + + this.drawLabel = function (ctx, t, x, y, w, h) { + ctx.fillStyle = "#aaa"; + ctx.font = "11px sans-serif"; + ctx.textAlign = "left"; + ctx.fillText(t, x, y + h / 2 + 4); + }; + + this.drawStringField = function (ctx, v, x, y, w, h) { + ctx.fillStyle = "#222"; + ctx.fillRect(x, y, w, h); + ctx.strokeStyle = "#444"; + ctx.strokeRect(x, y, w, h); + ctx.fillStyle = "#fff"; + ctx.font = "11px sans-serif"; + ctx.textAlign = "left"; + const d = v || ""; + ctx.fillText(d.length > 28 ? d.substring(0, 25) + "..." : d, x + 5, y + h / 2 + 4); + }; + + this.drawComboField = function (ctx, v, x, y, w, h) { + ctx.fillStyle = "#222"; + ctx.fillRect(x, y, w, h); + ctx.strokeStyle = "#444"; + ctx.strokeRect(x, y, w, h); + ctx.fillStyle = "#fff"; + ctx.font = "11px sans-serif"; + ctx.textAlign = "center"; + ctx.fillText(v, x + w / 2, y + h / 2 + 4); + ctx.fillStyle = "#666"; + ctx.beginPath(); + ctx.moveTo(x + w - 12, y + h / 2 - 3); + ctx.lineTo(x + w - 6, y + h / 2 - 3); + ctx.lineTo(x + w - 9, y + h / 2 + 3); + ctx.fill(); + }; + + this.onMouseDown = function (e, pos, canvas) { + for (const z of this.clickZones) { + if (pos[0] >= z.x && pos[0] <= z.x + z.w && + pos[1] >= z.y && pos[1] <= z.y + z.h) { + if (z.type === "prefix") { self.showPrefixInput(e); return true; } + if (z.type === "format") { self.showFormatSelector(e); return true; } + } + } + return false; + }; + + this.showPrefixInput = function (ev) { + const cv = self.rs_data.filename_prefix || ''; + const pop = document.createElement('div'); + pop.style.cssText = 'position:fixed;z-index:10002;background:#1a1a1a;border:1px solid #444;border-radius:6px;padding:8px 12px;box-shadow:0 4px 20px rgba(0,0,0,0.5);display:flex;align-items:center;gap:8px;'; + const inp = document.createElement('input'); + inp.type = 'text'; + inp.value = cv; + inp.style.cssText = 'width:220px;background:#222;color:#fff;border:1px solid #444;border-radius:4px;padding:6px 10px;font-size:12px;font-family:sans-serif;outline:none;'; + const btn = document.createElement('button'); + btn.textContent = 'OK'; + btn.style.cssText = 'background:#4CAF50;color:#fff;border:none;border-radius:4px;padding:6px 12px;font-size:12px;cursor:pointer;min-width:28px;'; + btn.onmouseover = () => btn.style.background = "#45a049"; + btn.onmouseout = () => btn.style.background = "#4CAF50"; + pop.appendChild(inp); + pop.appendChild(btn); + if (ev) { pop.style.left = (ev.clientX + 8) + 'px'; pop.style.top = (ev.clientY + 8) + 'px'; } + document.body.appendChild(pop); + setTimeout(() => { inp.focus(); if (cv.length) inp.select(); }, 50); + const save = () => { self.rs_data.filename_prefix = inp.value; self.syncToWidgets(); self.updateUI(); pop.remove(); }; + btn.onclick = (e) => { e.stopPropagation(); e.preventDefault(); save(); }; + inp.onkeydown = (e) => { if (e.key === 'Enter') { e.preventDefault(); save(); } }; + setTimeout(() => { + const cl = (e) => { if (!pop.contains(e.target)) { pop.remove(); document.removeEventListener("mousedown", cl); } }; + document.addEventListener("mousedown", cl); + }, 50); + }; + + this.showFormatSelector = function (ev) { + const FMTS = ["png", "jpg", "webp"]; + const menu = document.createElement("div"); + menu.style.cssText = 'position:fixed;background:#1a1a1a;border:1px solid #444;border-radius:6px;overflow:hidden;z-index:10001;box-shadow:0 4px 20px rgba(0,0,0,0.5);min-width:120px;'; + FMTS.forEach(f => { + const it = document.createElement("div"); + it.textContent = f.toUpperCase(); + it.style.cssText = 'padding:8px 15px;cursor:pointer;color:#ddd;font-size:12px;border-bottom:1px solid #333;'; + if (f === self.rs_data.format) { it.style.background = "#333"; it.style.color = "#4CAF50"; } + it.onmouseover = () => { if (f !== self.rs_data.format) it.style.background = "#333"; }; + it.onmouseout = () => { if (f !== self.rs_data.format) it.style.background = "#1a1a1a"; }; + it.onclick = (e) => { + e.stopPropagation(); e.preventDefault(); + self.rs_data.format = f; + self.syncToWidgets(); + self.updateUI(); + menu.remove(); + }; + menu.appendChild(it); + }); + if (ev) { menu.style.left = (ev.clientX + 8) + "px"; menu.style.top = (ev.clientY + 8) + "px"; } + document.body.appendChild(menu); + setTimeout(() => { + const cl = (e) => { if (!menu.contains(e.target)) { menu.remove(); document.removeEventListener("mousedown", cl); } }; + document.addEventListener("mousedown", cl); + }, 100); + }; + + this.updateUI = function () { + self.syncToWidgets(); + if (self.graph) self.graph.setDirtyCanvas(true, true); + }; + + const onSerialize = this.onSerialize; + this.onSerialize = function (o) { + self.syncToWidgets(); + return onSerialize ? onSerialize.apply(this, arguments) : undefined; + }; + + const onExecute = this.onExecute; + this.onExecute = function () { + self.syncToWidgets(); + return onExecute ? onExecute.apply(this, arguments) : undefined; + }; + + return result; + }; + } +}); \ No newline at end of file