Registry version change

README.md

@@ -4,15 +4,30 @@ Simple stack push/pop for images, strings and integers.
 The main use of this is to reduce unloading and loading of models when there are multiple steps in a workflow.
 E.g. using an SDXL model with ControlNet to create images to later feed into an Image to Video model.
 
-The Push nodes will save a file in the specified folder, default being a 'Stack' folder under the 'ComfyUI' folder. It will save a PNG for images, and text files for strings and numbers.
+## Push nodes
+
+The Push nodes will save a file in the specified folder, default being a 'Stack' folder under the 'ComfyUI' folder. It will save a PNG for images, and text files for strings and numbers. Generic objects are saved as pickle files.
+
+## Pop nodes
 
 The Pop will take the last file that was saved in the stack, and then delete it.
 The Pop nodes allow the option to override the pop, e.g. if you want to change the prompt for a generation without repushing another prompt. The stack will remain unchanged with the override enabled.
 The Pop nodes also allow a fallback when the stack is empty, without the fallback enabled the node with throw an exception stopping the generation. Note that any pops done before the exception in a generation will remain popped from the stacks.
 
+Popping now has four modes :  
+
+Normal - Normal stack behavior, last in first out  
+Peek - Peek at item on top of the stack, but don't remove  
+Pop Bottom - Reversed stack behavior, first in first out (i.e. the oldest file will be drawn)  
+Peek Bottom - Peek at the item at the bottom of the stack, don't remove  
+
+## Notes
+
 The 'Key' can be considered the variable name of the stack. For example, you could have two StackPushInt nodes, one with the key of 'width' and the other 'height', then be able to pull the right number in the next workflow for each variable.
 
-A generic object stack has also been implemented, but by default it is disabled.
+A generic object stack has also been implemented, but by default it is disabled. **Keeping this enabled and using other people's workflows does introduce a security risk.** I would advise you either keep this disabled when it is not needed, or clone this repo and change the node IDs to something unique.
+
+Suggestions/problems/improvements welcome.
 
 # Todo
 
@@ -22,6 +37,7 @@ A generic object stack has also been implemented, but by default it is disabled.
 
 ~~Add an override option~~
 
-Allow a 'peek' and 'sneak from bottom' option
+~~Allow a 'peek' and 'sneak from bottom' option~~
 
 Fix a 'no-change' prompt still having to be re-encoded
+- Needs investigation

__init__.py

@@ -5,7 +5,6 @@ import random
 import sys
 import traceback
 import shlex
-import os
 import numpy as np
 import torch
 import hashlib
@@ -28,6 +27,11 @@ def clean_filename(filename: str) -> str:
     
 def pil2tensor(image):
     return torch.from_numpy(np.array(image).astype(np.float32) / 255.0)
+
+def ensure_stack_dir(stackpath, key):
+    path = os.path.join(stackpath, clean_filename(key))
+    os.makedirs(path, exist_ok=True)
+    return path
  
 class AnyType(str):
     def __ne__(self, __value: object) -> bool:
@@ -72,15 +76,7 @@ class StackPushImage:
         if (key == ""):
             key = "default"
             
-        key = clean_filename(key)
-        
-        print("Pushing stack")
-        
-        path = stackpath + key + "/"
-
-        if not os.path.exists(path):
-            os.makedirs(path)
-            print(f"Created path: {path}")      
+        path = ensure_stack_dir(stackpath,key)  
         
         counter = 0
         
@@ -95,11 +91,11 @@ class StackPushImage:
                 if prompt is not None:
                     metadata.add_text("prompt", json.dumps(prompt))
 
-            filename = "stack" + now.strftime("%Y%m%d%H%M%S") + str(batch_number) + ".png"
+            filename = "stack" + now.strftime("%Y%m%d%H%M%S%f") + str(batch_number) + ".png"
             img.save(os.path.join(path, filename), pnginfo=metadata, compress_level=8)
             counter += 1
         
-        return { "ui": { "dunnohowtofixyet": "null" } }
+        return ( None, )
 
 class StackPushString:
     
@@ -134,24 +130,16 @@ class StackPushString:
         if (key == ""):
             key = "default"
         
-        key = clean_filename(key)
-        
-        print("Pushing stack")
-        
-        path = stackpath + key + "/"
-
-        if not os.path.exists(path):
-            os.makedirs(path)
-            print(f"Created path: {path}")      
-        
+        path = ensure_stack_dir(stackpath,key)  
+       
         #for (batch_number, stringv) in enumerate(stringvalue): 
         now = datetime.now()            
         
-        filename = os.path.join(path,"stack" + now.strftime("%Y%m%d%H%M%S") + "._st")
+        filename = os.path.join(path,"stack" + now.strftime("%Y%m%d%H%M%S%f") + "._st")
         with open(filename, 'w', encoding='utf-8') as f:
             f.write(stringvalue)
         
-        return { "ui": { "dunnohowtofixyet": "null" } }
+        return ( None, )
 
 class StackPushInt:
     
@@ -186,23 +174,15 @@ class StackPushInt:
         if (key == ""):
             key = "default"
         
-        key = clean_filename(key)
-        
-        print("Pushing stack")
-        
-        path = stackpath + key + "/"
-
-        if not os.path.exists(path):
-            os.makedirs(path)
-            print(f"Created path: {path}")      
+        path = ensure_stack_dir(stackpath,key)  
         
  #       for (batch_number, num) in enumerate(numbervalue):     
         now = datetime.now()
-        filename = os.path.join(path,"stack" + now.strftime("%Y%m%d%H%M%S") + "._nu")
+        filename = os.path.join(path,"stack" + now.strftime("%Y%m%d%H%M%S%f") + "._nu")
         with open(filename, 'w', encoding='utf-8') as f:
             f.write(str(numbervalue))
          
-        return { "ui": { "dunnohowtofixyet": "null" } }
+        return ( None, )
 
 class StackPushFloat:
     
@@ -237,23 +217,15 @@ class StackPushFloat:
         if (key == ""):
             key = "default"
         
-        key = clean_filename(key)
-        
-        print("Pushing stack")
-        
-        path = stackpath + key + "/"
-
-        if not os.path.exists(path):
-            os.makedirs(path)
-            print(f"Created path: {path}")      
+        path = ensure_stack_dir(stackpath,key)      
         
  #       for (batch_number, num) in enumerate(numbervalue):     
         now = datetime.now()
-        filename = os.path.join(path,"stack" + now.strftime("%Y%m%d%H%M%S") + "._fl")
+        filename = os.path.join(path,"stack" + now.strftime("%Y%m%d%H%M%S%f") + "._fl")
         with open(filename, 'w', encoding='utf-8') as f:
             f.write(str(numbervalue))
          
-        return { "ui": { "dunnohowtofixyet": "null" } }
+        return ( None, )
 
 class StackPushObject:
     
@@ -288,24 +260,16 @@ class StackPushObject:
         if (key == ""):
             key = "default"
         
-        key = clean_filename(key)
-        
-        print("Pushing stack")
-        
-        path = stackpath + key + "/"
-
-        if not os.path.exists(path):
-            os.makedirs(path)
-            print(f"Created path: {path}")      
+        path = ensure_stack_dir(stackpath,key)  
         
  #       for (batch_number, num) in enumerate(numbervalue):     
         now = datetime.now()
-        filename = os.path.join(path,"stack" + now.strftime("%Y%m%d%H%M%S") + "._pkl")
+        filename = os.path.join(path,"stack" + now.strftime("%Y%m%d%H%M%S%f") + "._pkl")
         
         with open(filename, 'wb') as f:
             pickle.dump(object, f)
          
-        return { "ui": { "dunnohowtofixyet": "null" } }
+        return ( None, )
 
                         
 
@@ -321,6 +285,7 @@ class StackPopImage:
                 "key": ("STRING", {"default": "default"}),
                 "fallback_enabled": ("BOOLEAN", {"default": False}),
                 "override_enabled": ("BOOLEAN", {"default": False}),
+                "stackmode": ( ['Normal', 'Peek', 'Pop Bottom', 'Peek Bottom'], {"default": "Normal"} ),
             },
             
             "optional": {
@@ -341,32 +306,35 @@ class StackPopImage:
     def IS_CHANGED(cls, **kwargs):
         return float("NaN")
    
-    def execute(self, key, fallback_enabled, override_enabled, fallback_item=None, override_item=None, stackpath="./ComfyUI/Stack/"):
+    def execute(self, key, fallback_enabled, override_enabled, stackmode, fallback_item=None, override_item=None, stackpath="./ComfyUI/Stack/"):
         
         if override_enabled:
             return (override_item,)
             
-        key = clean_filename(key)
-        
-        path = stackpath + key + "/"
+        path = os.path.join(stackpath,key)  
         
         images = []
         
         if not os.path.exists(path):
-            print(f"Key doesn't exist")   
             if (fallback_enabled):
                 return ( fallback_item, )
+            else:
+                raise Exception(f"Key doesn't exist") 
             
         # Get all .png files sorted by name
-        png_files = sorted([f for f in os.listdir(path) if f.lower().endswith('.png')] )
+        sorted_files = sorted([f for f in os.listdir(path) if f.lower().endswith('.png')] )
 
         last_file = ""
 
         # Check if there are any PNG files
-        if png_files:
-            last_file = os.path.join(path, png_files[-1])
+        if sorted_files:
+            if stackmode == 'Normal' or stackmode == 'Peek':
+                last_file = os.path.join(path, sorted_files[-1])
+            else:
+                last_file = os.path.join(path, sorted_files[0])
+
             image = Image.open(last_file)
-            print(f"Opened: {last_file}")
+            #print(f"Opened: {last_file}")
         else:
             if (fallback_enabled):
                 return ( fallback_item, )
@@ -379,7 +347,8 @@ class StackPopImage:
                 
         images.append(image)
         
-        os.remove(last_file)
+        if stackmode == 'Normal' or stackmode == 'Pop Bottom':
+            os.remove(last_file)
         
         return ( torch.cat(images, dim=0), )
 
@@ -395,6 +364,7 @@ class StackPopString:
                 "key": ("STRING", {"default": "default"}),
                 "fallback_enabled": ("BOOLEAN", {"default": False}),
                 "override_enabled": ("BOOLEAN", {"default": False}),
+                "stackmode": ( ['Normal', 'Peek', 'Pop Bottom', 'Peek Bottom'], {"default": "Normal"} ),
             },
             "optional": {
                 "stackpath": ("STRING",{"default": "./ComfyUI/Stack/"}),
@@ -414,42 +384,46 @@ class StackPopString:
     def IS_CHANGED(cls, **kwargs):
         return float("NaN")
    
-    def execute(self, key, fallback_enabled, fallback_item, override_enabled, override_item, stackpath="./ComfyUI/Stack/"):
+    def execute(self, key, fallback_enabled, fallback_item, override_enabled, override_item, stackmode, stackpath="./ComfyUI/Stack/"):
         
         if override_enabled:
             return (override_item,)
             
-        key = clean_filename(key)
-        
-        path = stackpath + key + "/"
+        path = os.path.join(stackpath,key)  
         
         images = []
         
         if not os.path.exists(path):
-            print(f"Key doesn't exist")   
             if (fallback_enabled):
                 return ( fallback_item, )
-            
-        # Get all .png files sorted by name
-        png_files = sorted([f for f in os.listdir(path) if f.lower().endswith('._st')] )
+            else:
+                raise Exception(f"Key doesn't exist") 
+                
+        # Get all ._st files sorted by name
+        sorted_files = sorted([f for f in os.listdir(path) if f.lower().endswith('._st')] )
 
         last_file = ""
 
         result = ""
 
-        # Check if there are any PNG files
-        if png_files:
-            last_file = os.path.join(path, png_files[-1])
+        # Check if there are any string files
+        if sorted_files:
+            if stackmode == 'Normal' or stackmode == 'Peek':
+                last_file = os.path.join(path, sorted_files[-1])
+            else:
+                last_file = os.path.join(path, sorted_files[0])
+
             with open(last_file, 'r', encoding='utf-8') as f:
                 result = f.read()
-            print(f"Opened: {last_file}")
+            #print(f"Opened: {last_file}")
         else:
             if (fallback_enabled):
                 return ( fallback_item, )
             else:
                 raise Exception("No string files found. Stack Empty.")
                      
-        os.remove(last_file)
+        if stackmode == 'Normal' or stackmode == 'Pop Bottom':
+            os.remove(last_file)
         
         return (result, )
 
@@ -465,6 +439,7 @@ class StackPopInt:
                 "key": ("STRING", {"default": "default"}),
                 "fallback_enabled": ("BOOLEAN", {"default": False}),
                 "override_enabled": ("BOOLEAN", {"default": False}),
+                "stackmode": ( ['Normal', 'Peek', 'Pop Bottom', 'Peek Bottom'], {"default": "Normal"} ),
             },
             "optional": {
                 "stackpath": ("STRING",{"default": "./ComfyUI/Stack/"}),
@@ -484,42 +459,46 @@ class StackPopInt:
     def IS_CHANGED(cls, **kwargs):
         return float("NaN")
    
-    def execute(self, key, fallback_enabled, fallback_item, override_enabled, override_item, stackpath="./ComfyUI/Stack/"):
+    def execute(self, key, fallback_enabled, fallback_item, override_enabled, override_item, stackmode, stackpath="./ComfyUI/Stack/"):
         
         if override_enabled:
             return (override_item,)
-            
-        key = clean_filename(key)
-        
-        path = stackpath + key + "/"
+
+        path = os.path.join(stackpath,key)  
         
         images = []
         
         if not os.path.exists(path):
-            print(f"Key doesn't exist")   
             if (fallback_enabled):
                 return ( fallback_item, )
-            
-        # Get all .png files sorted by name
-        png_files = sorted([f for f in os.listdir(path) if f.lower().endswith('._nu')] )
+            else:
+                raise Exception(f"Key doesn't exist") 
+                
+        # Get all ._nu files sorted by name
+        sorted_files = sorted([f for f in os.listdir(path) if f.lower().endswith('._nu')] )
 
         last_file = ""
 
         result = ""
 
-        # Check if there are any PNG files
-        if png_files:
-            last_file = os.path.join(path, png_files[-1])
+        # Check if there are any int files
+        if sorted_files:
+            if stackmode == 'Normal' or stackmode == 'Peek':
+                last_file = os.path.join(path, sorted_files[-1])
+            else:
+                last_file = os.path.join(path, sorted_files[0])
+                
             with open(last_file, 'r', encoding='utf-8') as f:
                 result = f.read()
-            print(f"Opened: {last_file}")
+            #print(f"Opened: {last_file}")
         else:
             if (fallback_enabled):
                 return ( fallback_item, )
             else:
                 raise Exception("No integer files found. Stack Empty.")
                       
-        os.remove(last_file)
+        if stackmode == 'Normal' or stackmode == 'Pop Bottom':
+            os.remove(last_file)
         
         return (int(result), )
 
@@ -535,6 +514,7 @@ class StackPopFloat:
                 "key": ("STRING", {"default": "default"}),
                 "fallback_enabled": ("BOOLEAN", {"default": False}),
                 "override_enabled": ("BOOLEAN", {"default": False}),
+                "stackmode": ( ['Normal', 'Peek', 'Pop Bottom', 'Peek Bottom'], {"default": "Normal"} ),
             },
             "optional": {
                 "stackpath": ("STRING",{"default": "./ComfyUI/Stack/"}),
@@ -554,42 +534,45 @@ class StackPopFloat:
     def IS_CHANGED(cls, **kwargs):
         return float("NaN")
    
-    def execute(self, key, fallback_enabled, fallback_item, override_enabled, override_item, stackpath="./ComfyUI/Stack/"):
+    def execute(self, key, fallback_enabled, fallback_item, override_enabled, stackmode, override_item, stackpath="./ComfyUI/Stack/"):
         
         if override_enabled:
             return (override_item,)
             
-        key = clean_filename(key)
-        
-        path = stackpath + key + "/"
+        path = os.path.join(stackpath,key)  
         
         images = []
         
         if not os.path.exists(path):
-            print(f"Key doesn't exist")   
             if (fallback_enabled):
                 return ( fallback_item, )
-            
-        # Get all .png files sorted by name
-        png_files = sorted([f for f in os.listdir(path) if f.lower().endswith('._fl')] )
+            else:
+                raise Exception(f"Key doesn't exist") 
+                
+        # Get all ._fl files sorted by name
+        sorted_files = sorted([f for f in os.listdir(path) if f.lower().endswith('._fl')] )
 
         last_file = ""
 
         result = ""
 
-        # Check if there are any PNG files
-        if png_files:
-            last_file = os.path.join(path, png_files[-1])
+        # Check if there are any float files
+        if sorted_files:
+            if stackmode == 'Normal' or stackmode == 'Peek':
+                last_file = os.path.join(path, sorted_files[-1])
+            else:
+                last_file = os.path.join(path, sorted_files[0])
             with open(last_file, 'r', encoding='utf-8') as f:
                 result = f.read()
-            print(f"Opened: {last_file}")
+            #print(f"Opened: {last_file}")
         else:
             if (fallback_enabled):
                 return ( fallback_item, )
             else:
                 raise Exception("No float files found. Stack Empty.")
                       
-        os.remove(last_file)
+        if stackmode == 'Normal' or stackmode == 'Pop Bottom':
+            os.remove(last_file)
         
         return (float(result), )
 
@@ -605,6 +588,7 @@ class StackPopObject:
                 "key": ("STRING", {"default": "default"}),
                 "fallback_enabled": ("BOOLEAN", {"default": False}),
                 "override_enabled": ("BOOLEAN", {"default": False}),
+                "stackmode": ( ['Normal', 'Peek', 'Pop Bottom', 'Peek Bottom'], {"default": "Normal"} ),
             },
             "optional": {
                 "stackpath": ("STRING",{"default": "./ComfyUI/Stack/"}),
@@ -626,7 +610,7 @@ class StackPopObject:
     def IS_CHANGED(cls, **kwargs):
         return float("NaN")
    
-    def execute(self, key, fallback_enabled, override_enabled, override_item=None, fallback_item=None, stackpath="./ComfyUI/Stack/"):
+    def execute(self, key, fallback_enabled, override_enabled, stackmode, override_item=None, fallback_item=None, stackpath="./ComfyUI/Stack/"):
         
         if not os.path.exists(os.path.join(folder_paths.get_user_directory(), "allow_generic_types.plz")):
             raise Exception("Generic object popping not allowed until dummy file \"allow_generic_types.plz\" exists in user folder ("+folder_paths.get_user_directory()+"). You should not do this unless you understand the risks of loading generic objects.")
@@ -634,37 +618,41 @@ class StackPopObject:
         if override_enabled:
             return (override_item,)
 
-        key = clean_filename(key)
-        
-        path = stackpath + key + "/"
+        path = os.path.join(stackpath,key)  
         
         images = []
         
         if not os.path.exists(path):
-            print(f"Key doesn't exist")   
             if (fallback_enabled):
                 return ( fallback_item, )
-            
-        # Get all .png files sorted by name
-        png_files = sorted([f for f in os.listdir(path) if f.lower().endswith('._pkl')] )
+            else:
+                raise Exception(f"Key doesn't exist") 
+                
+        # Get all .pkl files sorted by name
+        sorted_files = sorted([f for f in os.listdir(path) if f.lower().endswith('._pkl')] )
 
         last_file = ""
 
         result = ""
 
-        # Check if there are any PNG files
-        if png_files:
-            last_file = os.path.join(path, png_files[-1])
+        # Check if there are any pickle files
+        if sorted_files:
+            if stackmode == 'Normal' or stackmode == 'Peek':
+                last_file = os.path.join(path, sorted_files[-1])
+            else:
+                last_file = os.path.join(path, sorted_files[0])
+
             with open(last_file, 'rb') as f:
                 loaded_object = pickle.load(f)
-            print(f"Opened: {last_file}")
+            #print(f"Opened: {last_file}")
         else:
             if (fallback_enabled):
                 return ( fallback_item, )
             else:
                 raise Exception("No object files found. Stack Empty.")
                       
-        os.remove(last_file)
+        if stackmode == 'Normal' or stackmode == 'Pop Bottom':
+            os.remove(last_file)
         
         return (loaded_object,)#loaded_object,loaded_object,loaded_object )
 
@@ -688,7 +676,7 @@ NODE_DISPLAY_NAME_MAPPINGS = {
     "StackPushImage": "StackPushImage",
     "StackPushString": "StackPushString",
     "StackPushInt": "StackPushInt",
-    "StackPushInt": "StackPushInt",
+    "StackPushFloat": "StackPushFloat",
     "StackPushObject": "StackPushObject",
     "StackPopImage": "StackPopImage",
     "StackPopString": "StackPopString",

pyproject.toml

@@ -1,7 +1,7 @@
 [project]
 name = "ComfyUI-Stacker"
 description = "Simple stack push/pop style nodes for images, strings, integers and generic objects (image batches, latents, face models etc)"
-version = "0.1.2"
+version = "0.9"
 license = {file = "LICENSE"}
 dependencies = ["pickle"]