diff --git a/Loader.py b/Loader.py index 64000a9..8f8dfdf 100644 --- a/Loader.py +++ b/Loader.py @@ -72,6 +72,10 @@ class Loader: from .modules import Clamp modules.update(Clamp.NODE_CLASS_MAPPINGS) + if self.config()["modules"]["ImageBatch"]: + from .modules import ImageBatch + modules.update(ImageBatch.NODE_CLASS_MAPPINGS) + if self.config()["modules"]["ImageComposite"]: from .modules import ImageComposite modules.update(ImageComposite.NODE_CLASS_MAPPINGS) diff --git a/README.MD b/README.MD index 212a0ae..b997f2d 100644 --- a/README.MD +++ b/README.MD @@ -905,6 +905,53 @@ You also can change the fonts folder in config. +--- +### Image Batch +![image_batch](images/image_batch.png) +> Allows you to manage batch of images. + +
+ Nodes: + +### Get +![image_batch_get](images/image_batch_get.png) +> Get image from images batch. + +
+ Params: + +* index `[1 - *]` - Number of image. + +
+ +### Remove +![image_batch_remove](images/image_batch_remove.png) +> Remove image from images batch. + +
+ Params: + +* index `[1 - *]` - Number of image. + +
+ +### Fork +![image_batch_fork](images/image_batch_fork.png) +> Split images from batch. + +
+ Params: + +* priority `[first, second]` - If there are an odd number of images, choose which stream to send the image from the center to. + +
+ +### Join +![image_batch_join](images/image_batch_join.png) +> Combine images together into batch. + +
+ ## Configuration: diff --git a/config.json b/config.json index d52b709..0fcb9df 100644 --- a/config.json +++ b/config.json @@ -3,6 +3,7 @@ "modules": { "AlphaChanel": true, "Clamp": true, + "ImageBatch": true, "ImageComposite": true, "ImageContainer": true, "ImageDraw": true, diff --git a/modules/ImageBatch.py b/modules/ImageBatch.py new file mode 100644 index 0000000..37b64f1 --- /dev/null +++ b/modules/ImageBatch.py @@ -0,0 +1,136 @@ +import torch + + +class ImageBatchGet: + def __init__(self): + pass + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "images": ("IMAGE",), + "index": ("INT", { + "default": 1, + "min": 1, + "step": 1 + }), + }, + } + + RETURN_TYPES = ("IMAGE",) + FUNCTION = "image_batch_get" + CATEGORY = "image/batch" + + def image_batch_get(self, images, index): + batch = images.shape[0] + index = min(batch, index - 1) + + return (images[index].unsqueeze(0),) + + +class ImageBatchRemove: + def __init__(self): + pass + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "images": ("IMAGE",), + "index": ("INT", { + "default": 1, + "min": 1, + "step": 1 + }), + }, + } + + RETURN_TYPES = ("IMAGE",) + FUNCTION = "image_batch_get" + CATEGORY = "image/batch" + + def image_batch_get(self, images, index): + batch = images.shape[0] + index = min(batch, index - 1) + + return (torch.cat((images[:index], images[index + 1:]), dim=0),) + + +class ImageBatchFork: + def __init__(self): + pass + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "images": ("IMAGE",), + "priority": (["first", "second"],), + }, + } + + RETURN_TYPES = ("IMAGE", "IMAGE") + FUNCTION = "image_batch_join" + CATEGORY = "image/batch" + + def image_batch_join(self, images, priority): + batch = images.shape[0] + + if batch == 1: + return images, images + elif batch % 2 == 0: + first = batch // 2 + second = batch // 2 + else: + if priority == "first": + first = batch // 2 + 1 + second = batch // 2 + elif priority == "second": + first = batch // 2 + second = batch // 2 + 1 + else: + raise ValueError("Not existing priority.") + + return images[:first], images[-second:] + + +class ImageBatchJoin: + def __init__(self): + pass + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "images_a": ("IMAGE",), + "images_b": ("IMAGE",), + }, + } + + RETURN_TYPES = ("IMAGE",) + FUNCTION = "image_batch_join" + CATEGORY = "image/batch" + + def image_batch_join(self, images_a, images_b): + height_a, width_a, channels_a = images_a[0].shape + height_b, width_b, channels_b = images_b[0].shape + + if height_a != height_b: + raise ValueError("Height of images_a not equals of images_b. You can use ImageTransformResize for fix it.") + + if width_a != width_b: + raise ValueError("Width of images_a not equals of images_b. You can use ImageTransformResize for fix it.") + + if channels_a != channels_b: + raise ValueError("Channels of images_a not equals of images_b. Your can add or delete alpha channels with AlphaChanel module.") + + return (torch.cat((images_a, images_b)),) + + +NODE_CLASS_MAPPINGS = { + "ImageBatchGet": ImageBatchGet, + "ImageBatchRemove": ImageBatchRemove, + "ImageBatchFork": ImageBatchFork, + "ImageBatchJoin": ImageBatchJoin +}