diff --git a/web/fixed_comfyui_saver.py b/web/fixed_comfyui_saver.py deleted file mode 100644 index 9c54f91..0000000 --- a/web/fixed_comfyui_saver.py +++ /dev/null @@ -1,620 +0,0 @@ -import os -import json -import re -from datetime import datetime -from PIL import Image, PngImagePlugin -import numpy as np -import torch -import folder_paths -from PIL.PngImagePlugin import PngInfo -import platform - -class EndlessNode_Imagesaver: - """ - Enhanced batch image saver with comprehensive metadata support - Saves batched images with individual prompt names in filenames - Automatically handles multiple images from batch processing - Enhanced with workflow embedding, JSON export, and robust filename handling - """ - def __init__(self): - self.output_dir = folder_paths.get_output_directory() - self.type = "output" - self.compress_level = 4 - # OS-specific filename length limits - self.max_filename_length = self._get_max_filename_length() - - def _get_max_filename_length(self): - """Get maximum filename length based on OS""" - system = platform.system().lower() - if system == 'windows': - return 255 # NTFS limit - elif system in ['linux', 'darwin']: # Linux and macOS - return 255 # ext4/APFS limit - else: - return 200 # Conservative fallback - - @classmethod - def INPUT_TYPES(s): - return {"required": - {"images": ("IMAGE", ), - "prompt_list": ("STRING", {"forceInput": True}), - "include_timestamp": ("BOOLEAN", {"default": True}), - "timestamp_format": ("STRING", {"default": "%Y-%m-%d_%H-%M-%S", "description": "Use Python strftime format.\nExample: %Y-%m-%d %H-%M-%S\nSee: strftime.org for full options."}), - "image_format": (["PNG", "JPEG", "WEBP"], {"default": "PNG"}), - "jpeg_quality": ("INT", {"default": 95, "min": 1, "max": 100, "step": 1}), - "delimiter": ("STRING", {"default": "_"}), - "prompt_words_limit": ("INT", {"default": 8, "min": 1, "max": 16, "step": 1}), - "embed_workflow": ("BOOLEAN", {"default": True}), - "save_json_metadata": ("BOOLEAN", {"default": False}), - # ITEM #2: Enable/disable number padding - "enable_filename_numbering": ("BOOLEAN", {"default": True}), - # ITEM #1: Filename Number Padding Control - "filename_number_padding": ("INT", {"default": 2, "min": 1, "max": 9, "step": 1}), - "filename_number_start": ("BOOLEAN", {"default": False}), - # ITEM #3: Conditional PNG Metadata Embedding - "embed_png_metadata": ("BOOLEAN", {"default": True}), - }, - "optional": - {"output_path": ("STRING", {"default": ""}), - "filename_prefix": ("STRING", {"default": "Batch"}), - "negative_prompt_list": ("STRING", {"default": ""}), - "json_folder": ("STRING", {"default": ""}), - }, - "hidden": { - "prompt": "PROMPT", - "extra_pnginfo": "EXTRA_PNGINFO" - } - } - - RETURN_TYPES = ("STRING",) - RETURN_NAMES = ("saved_paths",) - FUNCTION = "save_batch_images" - OUTPUT_NODE = True - CATEGORY = "Endless 🌊✨/IO" - - def encode_emoji(self, obj): - """Properly encode emojis and special characters""" - if isinstance(obj, str): - return obj.encode('utf-8', 'surrogatepass').decode('utf-8') - return obj - - def clean_filename(self, text, max_words=8, delimiter="_"): - """Clean text for use in filenames with word limit and emoji support""" - # Limit to specified number of words - words = text.split()[:max_words] - text = ' '.join(words) - - # Handle emojis by encoding them properly - text = self.encode_emoji(text) - - # Replace illegal characters with delimiter, then clean up spaces - illegal_chars = r'[<>:"/\\|?*]' - clean_text = re.sub(illegal_chars, delimiter, text) - clean_text = re.sub(r'\s+', delimiter, clean_text) # Replace spaces with delimiter - clean_text = re.sub(r'[^\w\-_.{}]'.format(re.escape(delimiter)), '', clean_text) # Keep only safe chars - - return clean_text - - def format_timestamp(self, dt, format_string, delimiter='_'): - try: - formatted = dt.strftime(format_string) - # Replace colons first - formatted = formatted.replace(':', '-') - # Then replace all whitespace with the user's delimiter - if delimiter: - formatted = re.sub(r'\s+', delimiter, formatted) - return formatted - except Exception as e: - print(f"Invalid timestamp format: {e}") - return dt.strftime("%Y-%m-%d_%H-%M-%S") - - def validate_and_process_path(self, path, delimiter="_"): - if not path or path.strip() == "": - return path - - now = datetime.now() - - # Normalize path separators - path = path.replace("/", os.sep).replace("\\", os.sep) - - # Handle UNC or drive prefix - unc_prefix = "" - parts = path.split(os.sep) - - if path.startswith("\\\\"): # UNC path - if len(parts) >= 4: - unc_prefix = os.sep.join(parts[:4]) # \\server\share - parts = parts[4:] - else: - raise ValueError(f"Invalid UNC path: {path}") - elif re.match(r"^[A-Za-z]:$", parts[0]): # Drive letter - unc_prefix = parts[0] - parts = parts[1:] - - # Process the remaining subfolders - processed_parts = [] - for part in parts: - if not part: - continue - if "%" in part: - # Format date placeholders - formatted = self.format_timestamp(now, part, delimiter) - else: - # Sanitize folder names - formatted = re.sub(r'[<>:"/\\|?*]', delimiter, part) - processed_parts.append(formatted) - - # Reconstruct full path - full_path = os.path.join(unc_prefix, *processed_parts) - return full_path - - def ensure_filename_length(self, full_path, base_name, extension): - """Ensure the full filename doesn't exceed OS limits""" - directory = os.path.dirname(full_path) - - # Calculate available space for filename - dir_length = len(directory) + 1 # +1 for path separator - available_length = self.max_filename_length - len(extension) - max_base_length = available_length - dir_length - - if len(base_name) > max_base_length: - # Truncate base name to fit - base_name = base_name[:max_base_length-3] + "..." # -3 for ellipsis - - return os.path.join(directory, base_name + extension) - - def get_unique_filename(self, file_path): - """Generate unique filename by adding incremental numbers if file exists""" - if not os.path.exists(file_path): - return file_path - - directory = os.path.dirname(file_path) - filename = os.path.basename(file_path) - name, ext = os.path.splitext(filename) - - counter = 1 - while True: - new_name = f"{name}_{counter:03d}{ext}" - new_path = os.path.join(directory, new_name) - - # Check length constraints - if len(new_name) > self.max_filename_length: - # Truncate original name to make room for counter - available = self.max_filename_length - len(f"_{counter:03d}{ext}") - truncated_name = name[:available-3] + "..." - new_name = f"{truncated_name}_{counter:03d}{ext}" - new_path = os.path.join(directory, new_name) - - if not os.path.exists(new_path): - return new_path - counter += 1 - - def save_json_metadata(self, json_path, prompt_text, negative_text, - batch_index, creation_time, prompt=None, extra_pnginfo=None): - """Save JSON metadata file with ComfyUI-compatible format""" - - # Create the metadata structure that ComfyUI expects - metadata = { - "prompt": prompt_text, - "negative_prompt": negative_text, - "batch_index": batch_index, - "creation_time": creation_time, - } - - # BUG FIX #1: Save workflow and extra_pnginfo in the same format as PNG metadata - # ComfyUI expects these to be at the root level, not nested under other keys - if prompt is not None: - # The 'prompt' parameter contains the workflow data - metadata["workflow"] = prompt - - if extra_pnginfo is not None: - # extra_pnginfo contains additional workflow metadata - # Add each key-value pair from extra_pnginfo directly to metadata - for key, value in extra_pnginfo.items(): - metadata[key] = value - - # BUG FIX #2: Use the same JSON serialization as PNG metadata - # Don't double-encode the workflow data - try: - with open(json_path, 'w', encoding='utf-8', newline='\n') as f: - json.dump(metadata, f, indent=2, ensure_ascii=False) - return True - except Exception as e: - print(f"Failed to save JSON metadata: {e}") - return False - - def generate_numbered_filename(self, filename_prefix, delimiter, counter, - filename_number_padding, filename_number_start, - enable_filename_numbering, date_str, clean_prompt, ext): - """Generate filename with configurable number positioning and padding""" - # ITEM #3: Build filename parts in the correct order based on settings - filename_parts = [] - - # Always add timestamp first if provided - if date_str: - filename_parts.append(date_str) - - # Add number after timestamp if number_start is True AND numbering is enabled - if enable_filename_numbering and filename_number_start: - counter_str = f"{counter:0{filename_number_padding}}" - filename_parts.append(counter_str) - - # Add filename prefix if provided - if filename_prefix: - filename_parts.append(filename_prefix) - - # Add cleaned prompt - filename_parts.append(clean_prompt) - - # Add number at the end if number_start is False AND numbering is enabled - if enable_filename_numbering and not filename_number_start: - counter_str = f"{counter:0{filename_number_padding}}" - filename_parts.append(counter_str) - - # Join all parts with delimiter - filename = delimiter.join(filename_parts) + ext - - return filename - - def save_batch_images(self, images, prompt_list, include_timestamp=True, - timestamp_format="%Y-%m-%d_%H-%M-%S", image_format="PNG", - jpeg_quality=95, delimiter="_", - prompt_words_limit=8, embed_workflow=True, save_json_metadata=False, - enable_filename_numbering=True, filename_number_padding=2, - filename_number_start=False, embed_png_metadata=True, - output_path="", filename_prefix="batch", - negative_prompt_list="", json_folder="", prompt=None, extra_pnginfo=None): - - # Debug: Print tensor information - print(f"DEBUG: Images tensor shape: {images.shape}") - print(f"DEBUG: Images tensor type: {type(images)}") - - # BUG FIX #3: Debug workflow data structure - if save_json_metadata and prompt is not None: - print(f"DEBUG: Prompt type: {type(prompt)}") - print(f"DEBUG: Prompt keys: {list(prompt.keys()) if isinstance(prompt, dict) else 'Not a dict'}") - - if save_json_metadata and extra_pnginfo is not None: - print(f"DEBUG: Extra PNG info type: {type(extra_pnginfo)}") - print(f"DEBUG: Extra PNG info keys: {list(extra_pnginfo.keys()) if isinstance(extra_pnginfo, dict) else 'Not a dict'}") - - # Process output path with date/time validation (always process regardless of timestamp toggle) - processed_output_path = self.validate_and_process_path(output_path, delimiter) - - # Set output directory - if processed_output_path.strip() != "": - if not os.path.isabs(processed_output_path): - output_dir = os.path.join(self.output_dir, processed_output_path) - else: - output_dir = processed_output_path - else: - output_dir = self.output_dir - - # Create directory if it doesn't exist - try: - os.makedirs(output_dir, exist_ok=True) - except Exception as e: - raise ValueError(f"Could not create output directory {output_dir}: {e}") - - # Set up JSON directory - if save_json_metadata: - if json_folder.strip(): - processed_json_folder = self.validate_and_process_path(json_folder, delimiter) - if not os.path.isabs(processed_json_folder): - json_dir = os.path.join(self.output_dir, processed_json_folder) - else: - json_dir = processed_json_folder - else: - json_dir = output_dir - - try: - os.makedirs(json_dir, exist_ok=True) - except Exception as e: - print(f"Warning: Could not create JSON directory {json_dir}: {e}") - json_dir = output_dir - - # Generate datetime string if timestamp is enabled - now = datetime.now() - if include_timestamp: - date_str = self.format_timestamp(now, timestamp_format, delimiter) - else: - date_str = None - - # Parse individual prompts from the prompt list - individual_prompts = prompt_list.split('|') - individual_negatives = negative_prompt_list.split('|') if negative_prompt_list else [] - - # Set file extension - if image_format == "PNG": - ext = ".png" - elif image_format == "JPEG": - ext = ".jpg" - elif image_format == "WEBP": - ext = ".webp" - else: - ext = ".png" - - saved_paths = [] - - # Handle different tensor formats - if isinstance(images, torch.Tensor): - # Convert to numpy for easier handling - images_np = images.cpu().numpy() - print(f"DEBUG: Converted to numpy shape: {images_np.shape}") - - # Check if we have a batch dimension - if len(images_np.shape) == 4: # Batch format: [B, H, W, C] or [B, C, H, W] - batch_size = images_np.shape[0] - print(f"DEBUG: Found batch of {batch_size} images") - - for i in range(batch_size): - try: - # Extract single image from batch - img_array = images_np[i] - - # Validate and process the image array - if len(img_array.shape) != 3: - raise ValueError(f"Expected 3D tensor for image {i+1}, got shape {img_array.shape}") - - # Convert to 0-255 range if needed - if img_array.max() <= 1.0: - img_array = img_array * 255.0 - - img_array = np.clip(img_array, 0, 255).astype(np.uint8) - - # Handle different channel orders (HWC vs CHW) - if img_array.shape[0] == 3 or img_array.shape[0] == 4: # CHW format - img_array = np.transpose(img_array, (1, 2, 0)) # Convert to HWC - - img = Image.fromarray(img_array) - - # Get the corresponding prompt for this image - if i < len(individual_prompts): - prompt_text = individual_prompts[i].strip() - else: - # Cycle through prompts if we have more images than prompts - prompt_text = individual_prompts[i % len(individual_prompts)].strip() - print(f"Note: Cycling prompt for image {i+1} (using prompt {(i % len(individual_prompts)) + 1})") - - # Get corresponding negative prompt - negative_text = "" - if individual_negatives: - if i < len(individual_negatives): - negative_text = individual_negatives[i].strip() - else: - negative_text = individual_negatives[i % len(individual_negatives)].strip() - - # Clean the prompt for filename use - clean_prompt = self.clean_filename(prompt_text, prompt_words_limit, delimiter) - - # Generate filename using the new method - filename = self.generate_numbered_filename( - filename_prefix, delimiter, i+1, - filename_number_padding, filename_number_start, - enable_filename_numbering, date_str, clean_prompt, ext - ) - - # Create full file path and ensure length constraints - base_filename = os.path.splitext(filename)[0] - temp_path = os.path.join(output_dir, filename) - file_path = self.ensure_filename_length(temp_path, base_filename, ext) - - # Ensure unique filename - file_path = self.get_unique_filename(file_path) - - # Create JSON path if needed - if save_json_metadata: - json_base = os.path.splitext(os.path.basename(file_path))[0] - json_path = os.path.join(json_dir, json_base + ".json") - json_path = self.get_unique_filename(json_path) - - # Save image based on format - if image_format == "PNG": - # ITEM #3: Conditional PNG metadata embedding - if embed_png_metadata: - # Prepare PNG metadata - metadata = PngImagePlugin.PngInfo() - metadata.add_text("prompt", prompt_text) - metadata.add_text("negative_prompt", negative_text) - metadata.add_text("batch_index", str(i+1)) - metadata.add_text("creation_time", now.isoformat()) - - # Add workflow data if requested - if embed_workflow: - # BUG FIX #4: Save workflow data in the same format as expected by ComfyUI - if prompt is not None: - # Convert prompt to JSON string for PNG metadata - metadata.add_text("workflow", json.dumps(prompt, ensure_ascii=False)) - if extra_pnginfo is not None: - for key, value in extra_pnginfo.items(): - # Convert each value to JSON string for PNG metadata - metadata.add_text(key, json.dumps(value, ensure_ascii=False)) - - img.save(file_path, format="PNG", optimize=True, - compress_level=self.compress_level, pnginfo=metadata) - else: - # ITEM #3: Save clean PNG without metadata - img.save(file_path, format="PNG", optimize=True, - compress_level=self.compress_level) - - elif image_format == "JPEG": - # Convert RGBA to RGB for JPEG - if img.mode == 'RGBA': - background = Image.new('RGB', img.size, (255, 255, 255)) - background.paste(img, mask=img.split()[-1]) - img = background - img.save(file_path, format="JPEG", quality=jpeg_quality, optimize=True) - - elif image_format == "WEBP": - img.save(file_path, format="WEBP", quality=jpeg_quality, method=6) - - # Save JSON metadata if requested - if save_json_metadata: - self.save_json_metadata(json_path, prompt_text, negative_text, - i+1, now.isoformat(), prompt, extra_pnginfo) - - saved_paths.append(file_path) - print(f"Saved: {os.path.basename(file_path)}") - print(f" Prompt: {prompt_text}") - if negative_text: - print(f" Negative: {negative_text}") - if save_json_metadata: - print(f" JSON: {os.path.basename(json_path)}") - - except Exception as e: - error_msg = f"Failed to save image {i+1}: {e}" - print(error_msg) - # Continue with other images rather than failing completely - saved_paths.append(f"ERROR: {error_msg}") - - elif len(images_np.shape) == 3: # Single image format: [H, W, C] - print("DEBUG: Single image detected, processing as batch of 1") - # Process as single image - img_array = images_np - - # Convert to 0-255 range if needed - if img_array.max() <= 1.0: - img_array = img_array * 255.0 - - img_array = np.clip(img_array, 0, 255).astype(np.uint8) - img = Image.fromarray(img_array) - - prompt_text = individual_prompts[0].strip() if individual_prompts else "no_prompt" - negative_text = individual_negatives[0].strip() if individual_negatives else "" - clean_prompt = self.clean_filename(prompt_text, prompt_words_limit, delimiter) - - # Generate filename using the new method - filename = self.generate_numbered_filename( - filename_prefix, delimiter, 1, - filename_number_padding, filename_number_start, - enable_filename_numbering, date_str, clean_prompt, ext - ) - - base_filename = os.path.splitext(filename)[0] - temp_path = os.path.join(output_dir, filename) - file_path = self.ensure_filename_length(temp_path, base_filename, ext) - file_path = self.get_unique_filename(file_path) - - if save_json_metadata: - json_base = os.path.splitext(os.path.basename(file_path))[0] - json_path = os.path.join(json_dir, json_base + ".json") - json_path = self.get_unique_filename(json_path) - - if image_format == "PNG": - # ITEM #3: Conditional PNG metadata embedding - if embed_png_metadata: - metadata = PngImagePlugin.PngInfo() - metadata.add_text("prompt", prompt_text) - metadata.add_text("negative_prompt", negative_text) - metadata.add_text("batch_index", "1") - metadata.add_text("creation_time", now.isoformat()) - - if embed_workflow: - if prompt is not None: - metadata.add_text("workflow", json.dumps(prompt, ensure_ascii=False)) - if extra_pnginfo is not None: - for key, value in extra_pnginfo.items(): - metadata.add_text(key, json.dumps(value, ensure_ascii=False)) - - img.save(file_path, format="PNG", optimize=True, - compress_level=self.compress_level, pnginfo=metadata) - else: - # ITEM #3: Save clean PNG without metadata - img.save(file_path, format="PNG", optimize=True, - compress_level=self.compress_level) - elif image_format == "JPEG": - if img.mode == 'RGBA': - background = Image.new('RGB', img.size, (255, 255, 255)) - background.paste(img, mask=img.split()[-1]) - img = background - img.save(file_path, format="JPEG", quality=jpeg_quality, optimize=True) - elif image_format == "WEBP": - img.save(file_path, format="WEBP", quality=jpeg_quality, method=6) - - if save_json_metadata: - self.save_json_metadata(json_path, prompt_text, negative_text, - 1, now.isoformat(), prompt, extra_pnginfo) - - saved_paths.append(file_path) - print(f"Saved: {os.path.basename(file_path)}") - print(f" Prompt: {prompt_text}") - - else: - raise ValueError(f"Unexpected image tensor shape: {images_np.shape}") - - else: - # Handle case where images might be a list - print(f"DEBUG: Images is not a tensor, type: {type(images)}") - for i, image in enumerate(images): - try: - if isinstance(image, torch.Tensor): - img_array = image.cpu().numpy() - else: - img_array = np.array(image) - - # Process similar to above... - if img_array.max() <= 1.0: - img_array = img_array * 255.0 - - img_array = np.clip(img_array, 0, 255).astype(np.uint8) - - if len(img_array.shape) == 3 and (img_array.shape[0] == 3 or img_array.shape[0] == 4): - img_array = np.transpose(img_array, (1, 2, 0)) - - img = Image.fromarray(img_array) - - prompt_text = individual_prompts[i % len(individual_prompts)].strip() if individual_prompts else "no_prompt" - negative_text = individual_negatives[i % len(individual_negatives)].strip() if individual_negatives else "" - clean_prompt = self.clean_filename(prompt_text, prompt_words_limit, delimiter) - - # Generate filename using the new method - filename = self.generate_numbered_filename( - filename_prefix, delimiter, i+1, - filename_number_padding, filename_number_start, - enable_filename_numbering, date_str, clean_prompt, ext - ) - - base_filename = os.path.splitext(filename)[0] - temp_path = os.path.join(output_dir, filename) - file_path = self.ensure_filename_length(temp_path, base_filename, ext) - file_path = self.get_unique_filename(file_path) - - if save_json_metadata: - json_base = os.path.splitext(os.path.basename(file_path))[0] - json_path = os.path.join(json_dir, json_base + ".json") - json_path = self.get_unique_filename(json_path) - - # ITEM #3: Apply conditional PNG metadata for all image formats logic - if image_format == "PNG" and embed_png_metadata: - metadata = PngImagePlugin.PngInfo() - metadata.add_text("prompt", prompt_text) - metadata.add_text("negative_prompt", negative_text) - metadata.add_text("batch_index", str(i+1)) - metadata.add_text("creation_time", now.isoformat()) - - if embed_workflow: - if prompt is not None: - metadata.add_text("workflow", json.dumps(prompt, ensure_ascii=False)) - if extra_pnginfo is not None: - for key, value in extra_pnginfo.items(): - metadata.add_text(key, json.dumps(value, ensure_ascii=False)) - - img.save(file_path, format="PNG", optimize=True, - compress_level=self.compress_level, pnginfo=metadata) - else: - img.save(file_path, format=image_format.upper()) - - if save_json_metadata: - self.save_json_metadata(json_path, prompt_text, negative_text, - i+1, now.isoformat(), prompt, extra_pnginfo) - - saved_paths.append(file_path) - print(f"Saved: {os.path.basename(file_path)}") - - except Exception as e: - error_msg = f"Failed to save image {i+1}: {e}" - print(error_msg) - saved_paths.append(f"ERROR: {error_msg}") - - # Return all saved paths joined with newlines - return ("\n".join(saved_paths),) \ No newline at end of file