weekii-comfyui-save-image-pro/save_image_simple.py

"""
ComfyUI Save Image Pro - 简化版节点

简化版图像保存节点，提供基础的图像保存功能。

@version: 1.1.0
@author: AudioscavengeR
"""

import os
import sys
import json
import logging
from datetime import datetime
from typing import Dict, List, Any, Optional, Tuple
import folder_paths

# 导入重构后的核心模块
from .core import (
    SaveConfig, ConfigManager, ValidationResult,
    ImageSaver, FileNameGenerator, CounterManager,
    MetadataHandler, PathManager, JobDataExporter
)

# 版本信息
VERSION = "1.1.0"

# 设置日志
logger = logging.getLogger(__name__)

def check_optional_dependencies():
    """检查可选依赖项"""
    supported = {
        'avif': False,
        'jxl': False
    }
    
    try:
        from PIL import Image
        # 检查 AVIF 支持
        try:
            import pillow_avif
            supported['avif'] = True
        except ImportError:
            pass
        
        # 检查 JXL 支持
        try:
            import jxlpy
            supported['jxl'] = True
        except ImportError:
            pass
    except ImportError:
        pass
    
    return supported

# 检查支持的格式
SUPPORTED_FORMATS = check_optional_dependencies()


class SaveImageSimple:
    """
    ComfyUI Save Image Pro - 简化版
    
    提供基础的图像保存功能，包含最常用的参数。
    """
    
    RETURN_TYPES = ("IMAGE",)
    RETURN_NAMES = ("images",)
    FUNCTION = 'save_images'
    OUTPUT_NODE = True
    CATEGORY = 'image'
    DESCRIPTION = """
### ComfyUI Save Image Pro - 简化版

**基础功能:**
- 支持多种图像格式 (PNG, WebP, JPEG, AVIF, JXL, TIFF, GIF, BMP)
- 自定义文件名前缀
- 自定义文件夹前缀
- 质量控制
- 作业数据导出
- 图像预览

**简化设计:**
- 精简的参数设置
- 快速上手使用
- 保留核心功能
"""
    
    def __init__(self):
        self.config_manager = ConfigManager()
        self.output_dir = folder_paths.get_output_directory()
        self.image_saver = None
        self.job_exporter = None
        
        # 初始化日志
        self._setup_logging()
    
    def _setup_logging(self):
        """设置日志系统"""
        logging.basicConfig(
            level=logging.INFO,
            format='[%(name)s] %(levelname)s: %(message)s'
        )
        logger.info("ComfyUI Save Image Pro (Simple) initialized")
    
    @classmethod
    def INPUT_TYPES(cls):
        """定义输入类型 - 简化版"""
        # 获取支持的格式列表
        output_formats = ['.webp', '.png', '.jpg', '.jpeg', '.gif', '.tiff', '.bmp']
        
        if SUPPORTED_FORMATS['jxl']:
            output_formats.insert(0, '.jxl')
        if SUPPORTED_FORMATS['avif']:
            output_formats.insert(0, '.avif')
        
        return {
            "required": {
                "images": ("IMAGE", {"tooltip": "The images to save and pass through."}),
                "filename_prefix": ("STRING", {"default": "ComfyUI"}),
                "foldername_prefix": ("STRING", {"default": ""}),
                "output_format": (output_formats, {"default": ".webp"}),
                "quality": ("INT", {"default": 75, "min": 1, "max": 100, "step": 1}),
                "save_image_metadata": ("BOOLEAN", {"default": True, "tooltip": "Embed metadata in saved files. Turn off to save plain images."}),
                "metadata_mode": (["comfy", "a1111", "comfy+a1111", "off"], {"default": "comfy", "tooltip": "comfy writes prompt/workflow; a1111 writes share-site parameters; comfy+a1111 writes both."}),
                "sidecar_mode": (["off", "json", "txt", "both"], {"default": "off", "tooltip": "Optionally write metadata sidecar files next to each saved image."}),
                "image_preview": ("BOOLEAN", {"default": True}),
            },
            "hidden": {
                "prompt": "PROMPT",
                "extra_pnginfo": "EXTRA_PNGINFO"
            }
        }
    
    def save_images(self, images, filename_prefix="ComfyUI", foldername_prefix="",
                   output_format=".webp", quality=75, save_image_metadata=True,
                   metadata_mode="comfy", sidecar_mode="off", image_preview=True,
                   prompt=None, extra_pnginfo=None):
        """
        保存图像的主要方法 - 简化版
        """
        try:
            # 创建简化的配置对象
            config = SaveConfig(
                filename_prefix=filename_prefix,
                filename_keys="",  # 简化版不使用复杂的文件名键
                foldername_prefix=foldername_prefix,  # 使用用户输入的文件夹前缀
                foldername_keys="",  # 简化版不使用文件夹键
                delimiter="-",
                output_format=output_format,
                quality=quality,
                save_metadata=save_image_metadata,
                metadata_mode=metadata_mode,
                sidecar_mode=sidecar_mode,
                counter_digits=4,
                counter_position="last",
                one_counter_per_folder=True,
                save_job_data="disabled",  # 强制禁用JSON导出
                job_data_per_image=False, # 彻底禁用JSON导出
                job_custom_text="",
                image_preview=image_preview
            )
            
            # 验证配置
            validation = self.config_manager.validate_config(config)
            if not validation.valid:
                logger.error(f"Configuration validation failed: {validation.errors}")
                return {"ui": {"images": []}, "result": (images,)}
            
            # 初始化图像保存器
            if not self.image_saver:
                self.image_saver = ImageSaver(config, self.output_dir)
            else:
                self.image_saver.update_config(config)
            
            # 保存图像
            results = self.image_saver.save_images(
                images=images,
                prompt=prompt,
                extra_pnginfo=extra_pnginfo
            )
            
            # JSON导出功能已禁用 - 图片本身已包含完整元数据
            # 不再需要额外的JSON文件
            
            logger.info(f"Successfully saved {len(results)} images")
            
            # 返回结果
            if config.image_preview:
                return {"ui": {"images": results}, "result": (images,)}
            else:
                return {"ui": {"images": []}, "result": (images,)}
                
        except Exception as e:
            logger.error(f"Error occurred while saving images: {str(e)}")
            import traceback
            traceback.print_exc()
            return {"ui": {"images": []}, "result": (images,)}


# 节点类映射
NODE_CLASS_MAPPINGS = {
    "SaveImageSimple": SaveImageSimple
}

# 节点显示名称映射
NODE_DISPLAY_NAME_MAPPINGS = {
    "SaveImageSimple": "comfyui-save-image-pro (Simple)"
}