diff --git a/Rayko_VAE_Save_Image.py b/Rayko_VAE_Save_Image.py index 0beff54..1f5f5b0 100644 --- a/Rayko_VAE_Save_Image.py +++ b/Rayko_VAE_Save_Image.py @@ -17,16 +17,19 @@ import os import json import numpy as np import re -from datetime import datetime +import time from PIL import Image from PIL.PngImagePlugin import PngInfo import folder_paths import torch +from server import PromptServer +from aiohttp import web class RS_VAE_Decode_Save: def __init__(self): self.output_dir = folder_paths.get_output_directory() + self.last_temp_file = None @classmethod def INPUT_TYPES(cls): @@ -34,8 +37,10 @@ class RS_VAE_Decode_Save: "required": { "samples": ("LATENT",), "vae": ("VAE",), - "filename_prefix": ("STRING", {"default": "ComfyUI"}), + "save_path": ("STRING", {"default": ""}), + "file_prefix": ("STRING", {"default": "img"}), "format": (["png", "jpg", "webp"], {"default": "png"}), + "node_data": ("STRING", {"default": "{}", "multiline": False}), }, "hidden": { "prompt": "PROMPT", @@ -54,40 +59,44 @@ class RS_VAE_Decode_Save: JPG_QUALITY = 90 WEBP_QUALITY = 90 EMBED_WORKFLOW = True - OVERWRITE_EXISTING = False @staticmethod def _sanitize_path_component(component: str) -> str: - """Sanitizes a string to be safe for use in file paths.""" if not isinstance(component, str): component = str(component) component = component.replace("\\", "/") while "../" in component or "./" in component: component = component.replace("../", "").replace("./", "") - if os.path.isabs(component): - component = os.path.basename(component) component = re.sub(r'[<>:"|?*]', '_', component) - if not component.strip(): - component = "sanitized_empty" return component.strip() @staticmethod def _normalize_images(images): - while images.dim() > 4 and images.shape[0] == 1: - images = images.squeeze(0) + images = images.squeeze() if images.dim() == 2: images = images.unsqueeze(0).unsqueeze(-1) if images.dim() == 3: - images = images.unsqueeze(0) + if images.shape[-1] in [1, 3, 4]: + images = images.unsqueeze(0) + else: + images = images.unsqueeze(-1) + if images.dim() == 4: + b, h, w, c = images.shape + if c not in [1, 3, 4]: + if images.shape[1] in [1, 3, 4]: + images = images.permute(0, 2, 3, 1) + if images.dim() == 5: + if images.shape[1] == 1: + images = images.squeeze(1) + elif images.shape[0] == 1: + images = images.squeeze(0) + else: + images = images.reshape(images.shape[0], -1, images.shape[-2], images.shape[-1]) 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) + try: + images = images.view(images.shape[0], images.shape[-2], images.shape[-1], -1) + except: + raise ValueError(f"[RS] Cannot normalize tensor with shape: {images.shape}") if images.shape[-1] == 1: images = images.repeat(1, 1, 1, 3) if images.dtype != torch.float32: @@ -95,30 +104,9 @@ class RS_VAE_Decode_Save: images = torch.clamp(images, 0.0, 1.0) return images - def _get_next_filename(self, directory, base_name, extension, is_batch, batch_index=None): + def _get_next_counter(self, directory: str, prefix: str, extension: str) -> int: try: - safe_base_name = self._sanitize_path_component(base_name) - safe_extension = self._sanitize_path_component(extension).lower() - - test_path = os.path.join(directory, f"{safe_base_name}.{safe_extension}") - real_dir = os.path.realpath(directory) - real_path = os.path.realpath(test_path) - - if not real_path.startswith(real_dir + os.sep) and real_path != real_dir: - print(f"[RS] Security Warning: Path traversal detected in '{base_name}'. Falling back to safe name.") - safe_base_name = "unsafe_input_sanitized" - - if self.OVERWRITE_EXISTING: - if is_batch and batch_index is not None: - return os.path.join(directory, f"{safe_base_name}_{batch_index:03d}.{safe_extension}") - return os.path.join(directory, f"{safe_base_name}_001.{safe_extension}") - - if is_batch and batch_index is not None: - base_with_index = f"{safe_base_name}_{batch_index:03d}" - pattern = re.compile(rf'^{re.escape(base_with_index)}_(\d+)\.{re.escape(safe_extension)}$') - else: - pattern = re.compile(rf'^{re.escape(safe_base_name)}_(\d+)\.{re.escape(safe_extension)}$') - + pattern = re.compile(rf'^{re.escape(prefix)}_(\d{{5}})\.{re.escape(extension)}$') max_num = 0 if os.path.exists(directory): for f in os.listdir(directory): @@ -127,101 +115,75 @@ class RS_VAE_Decode_Save: num = int(match.group(1)) if num > max_num: max_num = num + return max_num + 1 + except Exception: + return 1 - next_num = max_num + 1 - if is_batch and batch_index is not None: - filename = f"{safe_base_name}_{batch_index:03d}_{next_num:03d}.{safe_extension}" - else: - filename = f"{safe_base_name}_{next_num:03d}.{safe_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"error_fallback_{ts}.png") - - def _build_workflow_json(self, prompt, extra_pnginfo): + def _cleanup_temp(self): 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"} + temp_dir = folder_paths.get_temp_directory() + if not os.path.exists(temp_dir): + return + + now = time.time() + cutoff = now - 7200 + + for filename in os.listdir(temp_dir): + if filename.startswith("rs_prev_"): + filepath = os.path.join(temp_dir, filename) + try: + if os.path.getmtime(filepath) < cutoff: + os.remove(filepath) + except: + pass except Exception as e: - print(f"[RS] Workflow JSON error: {e}") - return {"prompt": prompt or {}} + print(f"[RS] Cleanup error: {e}") - 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, + def decode_and_save(self, samples, vae, save_path, file_prefix, format, node_data, prompt=None, extra_pnginfo=None): images = vae.decode(samples["samples"]) images = self._normalize_images(images) - safe_prefix = self._sanitize_path_component(filename_prefix) + self._cleanup_temp() - full_output_folder, fn_prefix, counter, subfolder, _ = \ - folder_paths.get_save_image_path(safe_prefix, self.output_dir, - images[0].shape[1], images[0].shape[0]) + clean_path = save_path.strip() + is_absolute = os.path.isabs(clean_path) or (len(clean_path) > 1 and clean_path[1] == ':') + + if is_absolute: + target_dir = os.path.normpath(clean_path) + is_internal = False + else: + safe_path = self._sanitize_path_component(clean_path) + target_dir = os.path.join(self.output_dir, safe_path) if safe_path else self.output_dir + is_internal = True + + try: + os.makedirs(target_dir, exist_ok=True) + except Exception as e: + raise PermissionError(f"[RS] Cannot create directory: {target_dir}. Error: {e}") + + safe_prefix = self._sanitize_path_component(file_prefix) + if not safe_prefix: + safe_prefix = "img" batch_size = images.shape[0] - is_batch = batch_size > 1 ext_map = {"png": "png", "jpg": "jpg", "webp": "webp"} extension = ext_map.get(format, "png") + + start_counter = self._get_next_counter(target_dir, safe_prefix, extension) saved_files = [] + temp_dir = folder_paths.get_temp_directory() + + if self.last_temp_file and os.path.exists(os.path.join(temp_dir, self.last_temp_file)): + try: os.remove(os.path.join(temp_dir, self.last_temp_file)) + except: pass + 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) - + current_counter = start_counter + i + filename_main = f"{safe_prefix}_{current_counter:05}.{extension}" + filepath_main = os.path.join(target_dir, filename_main) + img_array = np.clip(255.0 * images[i].cpu().numpy(), 0, 255).astype(np.uint8) img = Image.fromarray(img_array) @@ -233,38 +195,66 @@ class RS_VAE_Decode_Save: 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 + try: metadata.add_text(key, json.dumps(value)) + except: 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_main, **save_kwargs) - img.save(filepath, **save_kwargs) + if is_internal: + rel_subfolder = os.path.relpath(target_dir, self.output_dir) + if rel_subfolder == ".": rel_subfolder = "" + + saved_files.append({ + "filename": filename_main, + "subfolder": rel_subfolder, + "type": "output" + }) + else: + temp_name = f"rs_prev_{int(time.time() * 1000)}_{i}.png" + filepath_temp = os.path.join(temp_dir, temp_name) + img.save(filepath_temp, compress_level=1) + + if i == batch_size - 1: + self.last_temp_file = temp_name - saved_files.append({ - "filename": filename, - "subfolder": subfolder, - "type": "output" - }) + saved_files.append({ + "filename": temp_name, + "subfolder": "", + "type": "temp" + }) return {"ui": {"images": saved_files}, "result": (images,)} +@PromptServer.instance.routes.get("/rs_folders") +async def get_output_folders(request): + try: + output_dir = folder_paths.get_output_directory() + subfolders = [] + if os.path.exists(output_dir): + for item in os.listdir(output_dir): + full_path = os.path.join(output_dir, item) + if os.path.isdir(full_path): + real_path = os.path.realpath(full_path) + real_output = os.path.realpath(output_dir) + if real_path.startswith(real_output + os.sep) or real_path == real_output: + subfolders.append(item) + return web.json_response({"subfolders": sorted(subfolders)}) + except Exception as e: + return web.json_response({"subfolders": [], "error": str(e)}, status=500) + + NODE_CLASS_MAPPINGS = { "RS_VAE_Decode_Save": RS_VAE_Decode_Save, } - NODE_DISPLAY_NAME_MAPPINGS = { - "RS_VAE_Decode_Save": "🦊 RS Decode Save", + "RS_VAE_Decode_Save": "🦊 RS Decode Save Image", } \ No newline at end of file