From 1640e399a876e57534896ffaafca6cf76c7cd7e4 Mon Sep 17 00:00:00 2001
From: Thomas Ward <teward@thomas-ward.net>
Date: Fri, 24 May 2024 21:54:33 -0400
Subject: [PATCH] Introduce advanced image saver

---
 _mappings.py |  4 +++-
 image.py     | 12 +++++-------
 2 files changed, 8 insertions(+), 8 deletions(-)

diff --git a/_mappings.py b/_mappings.py
index df9a5c6..a3bdd41 100644
--- a/_mappings.py
+++ b/_mappings.py
@@ -1,7 +1,7 @@
 from .generation import TWCUI_Util_GenerationParameters, TWCUI_Util_CommonSDXLResolutions
 from .generation import TWCUI_Util_GenerationPrompts
 
-from .image import TWCUI_Util_SaveImage
+from .image import TWCUI_Util_SaveImage, TWCUI_Util_SaveImageAdvanced
 
 from .util import TWCUI_Util_StringLiteral, TWCUI_Util_MultilineStringLiteral
 from .util import TWCUI_Util_IntLiteral, TWCUI_Util_FloatLiteral
@@ -14,6 +14,7 @@ NODE_CLASS_MAPPINGS = {
 
     # Image category
     "TWCUI_Util_SaveImage": TWCUI_Util_SaveImage,
+    "TWCUI_Util_SaveImageAdvanced": TWCUI_Util_SaveImageAdvanced,
 
     # Util category
     "TWCUI_Util_FloatLiteral": TWCUI_Util_FloatLiteral,
@@ -30,6 +31,7 @@ NODE_DISPLAY_NAME_MAPPINGS = {
 
     # Image category
     "TWCUI_Util_SaveImage": "Save Image",
+    "TWCUI_Util_SaveImageAdvanced": "Save Image (Advanced)",
 
     # Util category
     "TWCUI_Util_FloatLiteral": "FLOAT Literal",
diff --git a/image.py b/image.py
index 8b1d78b..779fcaa 100644
--- a/image.py
+++ b/image.py
@@ -112,9 +112,6 @@ class TWCUI_Util_SaveImageAdvanced(BaseNode):
     def __init__(self):
         super().__init__()
         self.output_dir = folder_paths.get_output_directory()
-        self.type = "output"
-        self.prefix_append = ""
-        self.png_compress_level = 4
 
     RETURN_TYPES = ()
 
@@ -153,6 +150,7 @@ class TWCUI_Util_SaveImageAdvanced(BaseNode):
                 "save_extra_pnginfo_with_metadata": ("BOOLEAN", {"default": False}),
                 "model_hash": ("STRING", {"default": "unknown", "forceInput": True}),
                 "vae_hash": ("STRING", {"default": "unknown", "forceInput": True}),
+                "compression": ("INT", {"default": 4, "min": 1, "max": 9, "step": 1}),
             },
             "hidden": {
                 "prompt": "PROMPT",
@@ -196,7 +194,7 @@ class TWCUI_Util_SaveImageAdvanced(BaseNode):
              model_name: str, vae_name: str, sampler_name: str, scheduler: str, positive_prompt: str,
              negative_prompt: str, seed_value: int, width: int, height: int, lossless_webp: bool,
              quality_jpeg_or_webp: str, time_format: str, save_metadata: bool, save_workflow_with_metadata: bool,
-             save_extra_pnginfo_with_metadata: bool, model_hash: str, vae_hash: str,
+             save_extra_pnginfo_with_metadata: bool, model_hash: str, vae_hash: str, compression: int,
              prompt: dict = None, extra_pnginfo: dict = None):
         if path or path == '':
             path = os.path.join(self.output_dir, path)
@@ -240,11 +238,11 @@ class TWCUI_Util_SaveImageAdvanced(BaseNode):
 
                 filename = f"{filename_prefix}.png"
                 file = os.path.join(self.output_dir, path, filename)
-                img.save(file, pnginfo=metadata, compress_level=self.compress_level)
+                img.save(file, pnginfo=metadata, compress_level=compression)
             else:
                 filename = f"{filename_prefix}.{extension}"
                 file = os.path.join(self.output_dir, path, filename)
-                img.save(file, compress_level=self.compress_level, quality_jpeg_or_webp=quality_jpeg_or_webp,
+                img.save(file, compress_level=compression, quality_jpeg_or_webp=quality_jpeg_or_webp,
                          lossless=lossless_webp,)
                 if save_metadata:
                     exif_bytes = piexif.dump({
@@ -259,4 +257,4 @@ class TWCUI_Util_SaveImageAdvanced(BaseNode):
 
         return {"ui": {"images": map(
             lambda fname: {"filename": fname, "subfolder": subfolder if subfolder != '.' else '',
-                           "type": 'output'}, paths)}}
\ No newline at end of file
+                           "type": 'output'}, paths)}}