General cleanup

.gitignore

@@ -157,4 +157,6 @@ cython_debug/
 #  be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore
 #  and can be added to the global gitignore or merged into this file.  For a more nuclear
 #  option (not recommended) you can uncomment the following to ignore the entire idea folder.
-.idea/
+.idea/
+
+storage/

Ostris_Nodes.py

@@ -1,4 +1,8 @@
-from .nodes.general import *
+from .utils.paths import fix_import_paths
+fix_import_paths()
+from .nodes.general import OstrisOneSeedNode, OstrisTextBoxNode
+from .nodes.batch_image_loader import OstrisBatchImageLoader
+
 
 ostris_node_list = [
     {
@@ -10,7 +14,12 @@ ostris_node_list = [
         "uid": "Text Box - Ostris",
         "class": OstrisTextBoxNode,
         "title": "Text Box",
-    }
+    },
+    # {
+    #     "uid": "Batch Image Loader - Ostris",
+    #     "class": OstrisBatchImageLoader,
+    #     "title": "Text Box",
+    # }
 ]
 
 # A dictionary that contains all nodes you want to export with their names

info.py

@@ -0,0 +1,7 @@
+INFO = {
+    "name": "Ostris Nodes",
+    "version": "0.0.1",
+    "author": "Ostris",
+    "project": "https://github.com/ostris/ostris_nodes_comfyui",
+    "description": "An extension for ComfyUI that adds Various QOL nodes",
+}

nodes/base_node.py

@@ -1,4 +1,8 @@
+from ..utils.storage import OstrisNodeStorage
+
 
 class OstrisBaseNode:
+    storage: OstrisNodeStorage('OstrisBaseNode')
+
     def __init__(self):
-        pass
+        pass

nodes/batch_image_loader.py

@@ -0,0 +1,163 @@
+import hashlib
+import os
+
+import numpy as np
+import torch
+
+from .base_node import OstrisBaseNode
+from ..settings.config import ostris_config
+from glob import glob
+from PIL import Image, ImageFilter, ImageEnhance, ImageOps, ImageDraw, ImageChops, ImageFont
+
+from ..utils.storage import OstrisNodeStorage
+
+
+def pil2tensor(image):
+    return torch.from_numpy(np.array(image).astype(np.float32) / 255.0).unsqueeze(0)
+
+
+def get_sha256(file_path):
+    sha256_hash = hashlib.sha256()
+    with open(file_path, 'rb') as file:
+        for chunk in iter(lambda: file.read(4096), b''):
+            sha256_hash.update(chunk)
+    return sha256_hash.hexdigest()
+
+
+OstrisBatchImageLoaderStorage = OstrisNodeStorage('OstrisBatchImageLoader')
+
+
+class ImgProcess:
+    to_process_image_path_list = []
+    processed_image_path_list = []
+    failed_image_path_list = []
+
+
+OstrisBatchImgProcess = ImgProcess()
+
+
+class OstrisBatchImageLoader(OstrisBaseNode):
+    last_folder_path = ''
+    img_process = OstrisBatchImgProcess
+
+    def __init__(self):
+        super().__init__()
+        self.storage = OstrisBatchImageLoaderStorage
+
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "folder_path": (
+                    "STRING", {
+                        "default": OstrisBatchImageLoaderStorage.get('last_folder_path', ''),
+                        "multiline": False
+                    }
+                ),
+            },
+        }
+
+    RETURN_TYPES = (
+        "IMAGE",
+        "STRING",
+        "STRING",
+        "STRING"
+        "STRING"
+        "STRING"
+    )
+    RETURN_NAMES = (
+        "image",
+        "filename",
+        "filename_no_ext",
+        "folder_path",
+        "full_path",
+        "caption"
+    )
+    FUNCTION = "batch_from_folder"
+
+    CATEGORY = f"{ostris_config.categories.batch}"
+
+    def load_image_paths(self, folder_path):
+        img_exts = ['jpg', 'jpeg', 'png', 'webp', 'tiff']
+        img_paths = []
+        files = glob(os.path.join(folder_path, '*'))
+        for file in files:
+            if file.split('.')[-1].lower() in img_exts:
+                img_paths.append(file)
+        self.img_process.to_process_image_path_list = img_paths
+        self.img_process.processed_image_path_list = []
+        self.img_process.failed_image_path_list = []
+
+    def get_caption(self, image_path):
+        # caption files will be named the same as the image without ext, but with a .txt or a .caption extension
+        caption_exts = ['txt', 'caption']
+        path_no_ext = os.path.splitext(image_path)[0]
+        possible_caption_paths = [f"{path_no_ext}.{ext}" for ext in caption_exts]
+
+        caption_path = None
+
+        for path in possible_caption_paths:
+            if os.path.exists(path):
+                caption_path = path
+                break
+
+        if caption_path is None:
+            return None
+        else:
+            with open(caption_path, 'r') as f:
+                raw = f.read()
+                arr = raw.split('\n')
+                clean = ", ".join([x.strip() for x in arr])
+                return clean
+
+    def batch_from_folder(self, **kwargs):
+        print('batch_from_folder', kwargs)
+        path = kwargs['folder_path']
+        self.storage.save('last_folder_path', path)
+
+        if not os.path.exists(path):
+            print('path does not exist')
+            raise Exception('Path does not exist')
+
+        if len(self.img_process.to_process_image_path_list) == 0:
+            print('No queued images, loading from folder')
+            # build our image paths list
+            self.load_image_paths(path)
+
+        image_path = self.img_process.to_process_image_path_list.pop(0)
+        try:
+            # process image
+            image = Image.open(image_path)
+            # flip based on meta
+            image = ImageOps.exif_transpose(image)
+            output_image = pil2tensor(image)
+            output_folder_path = os.path.dirname(image_path)
+            output_full_path = image_path
+            output_filename = os.path.basename(image_path)
+            output_filename_no_ext = os.path.splitext(output_filename)[0]
+            output_caption = self.get_caption(image_path)
+
+        except Exception as e:
+            print('OstrisBatchImageLoader ERROR:', e)
+            self.img_process.failed_image_path_list.append(image_path)
+            # try next image
+            return self.batch_from_folder(**kwargs)
+
+        self.img_process.processed_image_path_list.append(image_path)
+
+        return (
+            output_image,
+            output_filename,
+            output_filename_no_ext,
+            output_folder_path,
+            output_full_path,
+            output_caption
+        )
+
+    @classmethod
+    def IS_CHANGED(cls, **kwargs):
+        print('IS_CHANGED', kwargs)
+        print('cls.img_process', cls.img_process)
+        if len(cls.img_process.to_process_image_path_list):
+            return cls.img_process.to_process_image_path_list[0]
+        return False

nodes/general.py

@@ -77,11 +77,11 @@ class OstrisTextBoxNode(OstrisBaseNode):
 
     RETURN_TYPES = ("STRING",)
     RETURN_NAMES = ("text",)
-    FUNCTION = "text_multiline"
+    FUNCTION = "text_box"
 
     CATEGORY = ostris_config.categories.text
 
-    def text_multiline(self, text):
+    def text_box(self, text):
         import io
         new_text = []
         for line in io.StringIO(text):

settings/config.py

@@ -1,11 +1,29 @@
+import os
+from ..info import INFO
+
+OSTRIS_NODES_ROOT = os.path.dirname(os.path.dirname(__file__))
+OSTRIS_STORAGE_FOLDER = os.path.join(OSTRIS_NODES_ROOT, 'storage')
+OSTRIS_DB_FILE = os.path.join(OSTRIS_STORAGE_FOLDER, 'ostris.db.json')
+
+
+class OstrisPaths:
+    project_root = OSTRIS_NODES_ROOT
+    storage_folder = OSTRIS_STORAGE_FOLDER
+    db_file = OSTRIS_DB_FILE
+
+
 class OstrisCategory:
     root = "ostris"
     general = f"{root}/general"
     text = f"{root}/text"
+    image = f"{root}/image"
+    batch = f"{root}/batch"
 
 
 class OstrisConfig:
+    version = INFO['version']
     categories = OstrisCategory
+    paths = OstrisPaths
 
 
 ostris_config = OstrisConfig()

utils/paths.py

@@ -0,0 +1,11 @@
+
+def fix_import_paths():
+    # does all the import magical stuff
+    import os
+    import sys
+
+    sys.path.insert(0, os.path.join(os.path.dirname(os.path.realpath(__file__)), "ostris_nodes_comfyui"))
+    import folder_paths as comfy_paths
+
+    sys.path.append(comfy_paths.base_path)
+

utils/storage.py

@@ -0,0 +1,92 @@
+import json
+import os
+import typing
+from collections import OrderedDict
+
+from ..settings.config import ostris_config
+
+base_db = OrderedDict({
+    'version': ostris_config.version,
+    'nodes': {},
+})
+
+
+class OstrisDB:
+    _db: OrderedDict
+
+    def __init__(self):
+        self.db_path = ostris_config.paths.db_file
+        # make sure paths exist
+        if not os.path.exists(ostris_config.paths.storage_folder):
+            os.makedirs(ostris_config.paths.storage_folder, exist_ok=True)
+        if not os.path.exists(ostris_config.paths.db_file):
+            with open(ostris_config.paths.db_file, 'w') as f:
+                f.write(json.dumps(base_db, indent=4))
+
+        self._load_db()
+        self._migrate_db()
+
+    def _save_db(self):
+        with open(self.db_path, 'w') as f:
+            f.write(json.dumps(self._db, indent=4))
+
+    def _migrate_db(self):
+        data = self._db
+        did_update = False
+        for key, value in base_db.items():
+            if key not in data:
+                data[key] = value
+                did_update = True
+            if key == 'version':
+                data[key] = value
+                did_update = True
+        self._db = data
+        if did_update:
+            self._save_db()
+        return data, did_update
+
+    def _load_db(self):
+        with open(self.db_path, 'r') as f:
+            self._db = json.loads(f.read(), object_pairs_hook=OrderedDict)
+
+    def save_node_data(self, node_id, key, value):
+        if node_id not in self._db:
+            self._db[node_id] = {}
+        self._db[node_id][key] = value
+        self._save_db()
+
+    def get_node_data(self, node_id, key, default=None):
+        if node_id not in self._db:
+            return default
+        if key not in self._db[node_id]:
+            return default
+        return self._db[node_id][key]
+
+
+# only keep one instance of this to keep them all in sync
+_shared_db = None
+
+# prevent recursive import
+if typing.TYPE_CHECKING:
+    from ..nodes.base_node import OstrisBaseNode
+
+
+class OstrisNodeStorage:
+    db: 'OstrisDB'
+
+    def __init__(
+            self,
+            node_class_name: str,
+    ):
+        self._node_class_name_ = node_class_name
+
+        global _shared_db
+        if _shared_db is None:
+            _shared_db = OstrisDB()
+        self._db = _shared_db
+
+    def save(self, key: str, value):
+        self._db.save_node_data(self._node_class_name_, key, value)
+
+    def get(self, key: str, default=None):
+        return self._db.get_node_data(self._node_class_name_, key, default)