Moved batch images to sn0w utility

src/image_batch.py

@@ -1,44 +0,0 @@
-import torch
-
-class WAS_Image_Batch:
-    def __init__(self):
-        pass
-
-    @classmethod
-    def INPUT_TYPES(cls):
-        return {
-            "required": {
-            },
-            "optional": {
-                "images_a": ("IMAGE",),
-                "images_b": ("IMAGE",),
-                "images_c": ("IMAGE",),
-                "images_d": ("IMAGE",),
-                # "images_e": ("IMAGE",),
-                # "images_f": ("IMAGE",),
-                # Theoretically, an infinite number of image input parameters can be added.
-            },
-        }
-
-    RETURN_TYPES = ("IMAGE",)
-    RETURN_NAMES = ("image",)
-    FUNCTION = "image_batch"
-    CATEGORY = "WAS Suite/Image"
-
-    def _check_image_dimensions(self, tensors, names):
-        reference_dimensions = tensors[0].shape[1:]  # Ignore batch dimension
-        mismatched_images = [names[i] for i, tensor in enumerate(tensors) if tensor.shape[1:] != reference_dimensions]
-
-        if mismatched_images:
-            raise ValueError(f"WAS Image Batch Warning: Input image dimensions do not match for images: {mismatched_images}")
-
-    def image_batch(self, **kwargs):
-        batched_tensors = [kwargs[key] for key in kwargs if kwargs[key] is not None]
-        image_names = [key for key in kwargs if kwargs[key] is not None]
-
-        if not batched_tensors:
-            raise ValueError("At least one input image must be provided.")
-
-        self._check_image_dimensions(batched_tensors, image_names)
-        batched_tensors = torch.cat(batched_tensors, dim=0)
-        return (batched_tensors,)

src/lora_tester.py

@@ -1,6 +1,6 @@
 from nodes import KSampler, KSamplerAdvanced, VAEDecode, VAEEncode, EmptyLatentImage, CLIPTextEncode
 from custom_nodes.comfyui_lora_tag_loader.nodes import LoraTagLoader
-from .image_batch import WAS_Image_Batch
+from .sn0w import Utility
 from comfy_extras.nodes_upscale_model import ImageUpscaleWithModel
 import comfy.samplers
 
@@ -40,7 +40,7 @@ class LoraTestNode:
         vae_encode = VAEEncode()
         text_encode = CLIPTextEncode()
         lora_loader = LoraTagLoader()
-        image_batcher = WAS_Image_Batch()
+        utility = Utility()
         upscaler = ImageUpscaleWithModel()
 
         latent_image = EmptyLatentImage().generate(width, height)[0]
@@ -74,5 +74,6 @@ class LoraTestNode:
         image_batch_kwargs = {f"images_{chr(97 + i)}": image for i, image in enumerate(images)}
 
         # Using WAS_Image_Batch to batch the images together
-        batched_images = image_batcher.image_batch(**image_batch_kwargs)
+        batched_images = utility.image_batch(**image_batch_kwargs)
+        
         return batched_images

src/lora_tester_xl.py

@@ -1,5 +1,5 @@
 from nodes import KSampler, KSamplerAdvanced, VAEDecode, VAEEncode, EmptyLatentImage, LoraLoader, ImageScaleBy
-from .image_batch import WAS_Image_Batch
+from .sn0w import Utility
 from comfy_extras.nodes_clip_sdxl import CLIPTextEncodeSDXL
 from comfy_extras.nodes_upscale_model import ImageUpscaleWithModel
 import comfy.samplers
@@ -41,7 +41,6 @@ class LoraTestXLNode:
         vae_encode = VAEEncode()
         text_encode_xl = CLIPTextEncodeSDXL()
         lora_loader = LoraLoader()
-        image_batcher = WAS_Image_Batch()
         upscaler = ImageUpscaleWithModel()
         scale_image = ImageScaleBy()
 
@@ -104,6 +103,7 @@ class LoraTestXLNode:
         # Dynamically construct kwargs for image_batch
         image_batch_kwargs = {f"images_{chr(97 + i)}": image for i, image in enumerate(images)}
 
-        # Using WAS_Image_Batch to batch the images together
-        batched_images = image_batcher.image_batch(**image_batch_kwargs)
-        return (batched_images[0], len(images))
+        # Batch the images together
+        batched_images = Utility.image_batch(**image_batch_kwargs)
+
+        return (batched_images, len(images))

src/sn0w.py

@@ -1,5 +1,6 @@
 import os
 import json
+import torch
 
 class ConfigReader:
     _config = None
@@ -63,4 +64,24 @@ class Utility:
                 current_row.append(min(insertions, deletions, substitutions))
             previous_row = current_row
 
-        return previous_row[-1]
+        return previous_row[-1]
+    
+    @staticmethod
+    def _check_image_dimensions(tensors, names):
+        reference_dimensions = tensors[0].shape[1:]  # Ignore batch dimension
+        mismatched_images = [names[i] for i, tensor in enumerate(tensors) if tensor.shape[1:] != reference_dimensions]
+
+        if mismatched_images:
+            raise ValueError(f"Input image dimensions do not match for images: {mismatched_images}")
+
+    @staticmethod
+    def image_batch(**kwargs):
+        batched_tensors = [kwargs[key] for key in kwargs if kwargs[key] is not None]
+        image_names = [key for key in kwargs if kwargs[key] is not None]
+
+        if not batched_tensors:
+            raise ValueError("At least one input image must be provided.")
+
+        Utility._check_image_dimensions(batched_tensors, image_names)
+        batched_tensors = torch.cat(batched_tensors, dim=0)
+        return batched_tensors