From c1cb2e040b8bedc5955dbcde229f36f38d46d0e0 Mon Sep 17 00:00:00 2001
From: Fill <55672949+filliptm@users.noreply.github.com>
Date: Mon, 15 Jul 2024 23:36:30 -0700
Subject: [PATCH] Add files via upload

---
 __init__.py               |  3 +++
 fl_image_caption_saver.py | 55 +++++++++++++++++++++++++++++++++++++++
 2 files changed, 58 insertions(+)
 create mode 100644 fl_image_caption_saver.py

diff --git a/__init__.py b/__init__.py
index 670cd5d..dc31b9c 100644
--- a/__init__.py
+++ b/__init__.py
@@ -5,6 +5,7 @@ from .FL_KohyaSSDatasetConfig import FL_KohyaSSDatasetConfig
 from .FL_KohyaSSAdvConfig import FL_KohyaSSAdvConfig
 from .FL_KohyaSSTrain import FL_KohyaSSTrain
 from .FL_LoadImagesFromDirectoryPath import FL_LoadImagesFromDirectoryPath
+from .fl_image_caption_saver import FL_ImageCaptionSaver
 #==============================================================================
 #==============================================================================
 # from .FL_SliderLoraInitWorkspace import FL_SliderLoraInitWorkspace
@@ -19,6 +20,7 @@ NODE_CLASS_MAPPINGS = {
     "FL_KohyaSSAdvConfig": FL_KohyaSSAdvConfig,
     "FL_KohyaSSTrain": FL_KohyaSSTrain,
     "FL_LoadImagesFromDirectoryPath": FL_LoadImagesFromDirectoryPath,
+    "FL_ImageCaptionSaver": FL_ImageCaptionSaver,
 #==============================================================================
 #==============================================================================
     # "FL_SliderLoraInitWorkspace": FL_SliderLoraInitWorkspace,
@@ -33,6 +35,7 @@ NODE_DISPLAY_NAME_MAPPINGS = {
     "FL_KohyaSSAdvConfig": "FL Kohya Adv Config",
     "FL_KohyaSSTrain": "FL Kohya Train",
     "FL_LoadImagesFromDirectoryPath": "FL Kohya Data Loader",
+    "FL_ImageCaptionSaver": "FL Image Caption Saver"
 #==============================================================================
 #==============================================================================
     # "FL_SliderLoraInitWorkspace": "FL Slider LoRA Init Workspace",
diff --git a/fl_image_caption_saver.py b/fl_image_caption_saver.py
new file mode 100644
index 0000000..06c86cc
--- /dev/null
+++ b/fl_image_caption_saver.py
@@ -0,0 +1,55 @@
+import os
+from PIL import Image
+
+from comfy.utils import ProgressBar
+
+class FL_ImageCaptionSaver:
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "images": ("IMAGE", {}),
+                "folder_name": ("STRING", {"default": "output_folder"}),
+                "caption_text": ("STRING", {"default": "Your caption here"}),
+                "overwrite": ("BOOLEAN", {"default": True})  # New overwrite toggle
+            }
+        }
+
+    RETURN_TYPES = ("STRING",)
+    FUNCTION = "save_images_with_captions"
+    CATEGORY = "🏵️Fill Nodes/utility"
+    OUTPUT_NODE = True
+
+    def save_images_with_captions(self, images, folder_name, caption_text, overwrite):
+        # Ensure output directory exists
+        os.makedirs(folder_name, exist_ok=True)
+
+        saved_files = []
+        pbar = ProgressBar(len(images))
+        for i, image_tensor in enumerate(images):
+            base_name = f"image_{i}"
+            image_file_name = f"{folder_name}/{base_name}.png"
+            text_file_name = f"{folder_name}/{base_name}.txt"
+
+            # Check if overwrite is disabled and file exists
+            if not overwrite:
+                counter = 1
+                while os.path.exists(image_file_name) or os.path.exists(text_file_name):
+                    image_file_name = f"{folder_name}/{base_name}_{counter}.png"
+                    text_file_name = f"{folder_name}/{base_name}_{counter}.txt"
+                    counter += 1
+
+            # Convert tensor to image
+            image = Image.fromarray((image_tensor.numpy() * 255).astype('uint8'), 'RGB')
+
+            # Save image
+            image.save(image_file_name)
+            saved_files.append(image_file_name)
+
+            # Save text file
+            with open(text_file_name, "w") as text_file:
+                text_file.write(caption_text)
+
+            pbar.update_absolute(i)
+
+        return (f"Saved {len(images)} images and captions in '{folder_name}'",)