From 099d46e3554e9ba8dc6d9e2eda0794ebf72d4b31 Mon Sep 17 00:00:00 2001 From: sn0w12 Date: Fri, 23 Feb 2024 17:09:47 +0100 Subject: [PATCH] Moved batch images to sn0w utility --- src/image_batch.py | 44 ------------------------------------------- src/lora_tester.py | 7 ++++--- src/lora_tester_xl.py | 10 +++++----- src/sn0w.py | 23 +++++++++++++++++++++- 4 files changed, 31 insertions(+), 53 deletions(-) delete mode 100644 src/image_batch.py diff --git a/src/image_batch.py b/src/image_batch.py deleted file mode 100644 index c965d4c..0000000 --- a/src/image_batch.py +++ /dev/null @@ -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,) \ No newline at end of file diff --git a/src/lora_tester.py b/src/lora_tester.py index c3fee5e..6446f00 100644 --- a/src/lora_tester.py +++ b/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 \ No newline at end of file diff --git a/src/lora_tester_xl.py b/src/lora_tester_xl.py index 7eea3d2..a617a53 100644 --- a/src/lora_tester_xl.py +++ b/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)) \ No newline at end of file + # Batch the images together + batched_images = Utility.image_batch(**image_batch_kwargs) + + return (batched_images, len(images)) \ No newline at end of file diff --git a/src/sn0w.py b/src/sn0w.py index 32c98b8..acb4bd1 100644 --- a/src/sn0w.py +++ b/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] \ No newline at end of file + 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 \ No newline at end of file