259 lines
7.6 KiB
Python
259 lines
7.6 KiB
Python
import torch
|
|
import numpy as np
|
|
|
|
from . import register_node
|
|
|
|
|
|
@register_node
|
|
class JoinImageBatch:
|
|
"""Turns an image batch into one big image."""
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"images": ("IMAGE",),
|
|
"mode": (("horizontal", "vertical"), {"default": "horizontal"}),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE",)
|
|
RETURN_NAMES = ("image",)
|
|
INPUT_IS_LIST = True
|
|
FUNCTION = "join"
|
|
CATEGORY = "🔁 Sequential Batcher/Image"
|
|
|
|
def join(self, images, mode):
|
|
# mode will be a list of strings, e.g. ["horizontal", "horizontal", ...]
|
|
mode = mode[0]
|
|
if isinstance(images, list):
|
|
# images is a list of [1, H, W, C] tensors or similar
|
|
images = torch.cat(images, dim=0)
|
|
n, h, w, c = images.shape
|
|
image = None
|
|
if mode == "vertical":
|
|
# for vertical we can just reshape
|
|
image = images.reshape(1, n * h, w, c)
|
|
elif mode == "horizontal":
|
|
# for horizontal we have to swap axes
|
|
image = torch.transpose(torch.transpose(images, 1, 2).reshape(1, n * w, h, c), 1, 2)
|
|
return (image,)
|
|
|
|
|
|
@register_node
|
|
class JoinImages:
|
|
"""Turns joins two images into one big image."""
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"image_a": ("IMAGE",),
|
|
"image_b": ("IMAGE",),
|
|
"mode": (("horizontal", "vertical"), {"default": "horizontal"}),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE",)
|
|
RETURN_NAMES = ("image",)
|
|
FUNCTION = "join"
|
|
CATEGORY = "🔁 Sequential Batcher/Image"
|
|
|
|
def join(self, image_a, image_b, mode):
|
|
dim = {'horizontal': 2, 'vertical': 1}[mode]
|
|
return (torch.concat((image_a, image_b), dim), )
|
|
|
|
|
|
@register_node
|
|
class SelectImageBatch:
|
|
"""Selects one image from an image batch."""
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"images": ("IMAGE",),
|
|
"select": ("INT", {"default": 0, "min": 0, "max": 99999, "step": 1}),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE",)
|
|
RETURN_NAMES = ("image",)
|
|
FUNCTION = "select"
|
|
CATEGORY = "🔁 Sequential Batcher/Image"
|
|
|
|
def select(self, images, select):
|
|
n, h, w, c = images.shape
|
|
if select >= n:
|
|
select = n - 1
|
|
return (images[select].reshape(1, h, w, c),)
|
|
|
|
|
|
@register_node
|
|
class SelectImageList:
|
|
"""Selects one image from an image list."""
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"images": ("IMAGE",),
|
|
"select": ("INT", {"default": 0, "min": 0, "max": 99999, "step": 1}),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE",)
|
|
RETURN_NAMES = ("image",)
|
|
INPUT_IS_LIST = True
|
|
FUNCTION = "select"
|
|
CATEGORY = "🔁 Sequential Batcher/Image"
|
|
|
|
def select(self, images, select):
|
|
select = select[0]
|
|
n = len(images)
|
|
if select >= n:
|
|
select = n - 1
|
|
return (images[select],)
|
|
|
|
|
|
@register_node
|
|
class GetImageSize:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"images": ("IMAGE",),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("INT", "INT", "INT")
|
|
RETURN_NAMES = ("width", "height", "batch_size")
|
|
FUNCTION = "go"
|
|
CATEGORY = "🔁 Sequential Batcher/Image"
|
|
|
|
def go(self, images):
|
|
return (images.shape[2], images.shape[1], images.shape[0])
|
|
|
|
|
|
@register_node
|
|
class StringToImage:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"text": ("STRING", {"default": "Hello world!"}),
|
|
"width": ("INT", {"default": 384}),
|
|
"height": ("INT", {"default": 16}),
|
|
"colour": ("COLOR", {"default": "white"}),
|
|
"background": ("COLOR", {"default": "black"}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE",)
|
|
RETURN_NAMES = ("image",)
|
|
FUNCTION = "render"
|
|
CATEGORY = "🔁 Sequential Batcher/Image"
|
|
|
|
def render(self, text, width, height, colour, background):
|
|
from PIL import Image, ImageDraw, ImageFont
|
|
font = ImageFont.load_default()
|
|
img = Image.new("RGB", (width, height), background)
|
|
draw = ImageDraw.Draw(img)
|
|
_, _, w, h = draw.textbbox((0, 0), text, font=font)
|
|
draw.text(((width - w) / 2, ((height - h) / 2) - 1), text, font=font, fill=colour)
|
|
tensor = torch.from_numpy(np.array(img).astype(np.float32) / 255.0).unsqueeze(0)
|
|
return (tensor,)
|
|
|
|
|
|
@register_node
|
|
class ProgressBar:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"progress": ("FLOAT", {'default': 0, 'forceInput': True}),
|
|
"padding": ("INT", {'default': 3}),
|
|
"width": ("INT", {"default": 384}),
|
|
"height": ("INT", {"default": 16}),
|
|
"colour": ("COLOR", {"default": "white"}),
|
|
"background": ("COLOR", {"default": "black"}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE",)
|
|
RETURN_NAMES = ("image",)
|
|
FUNCTION = "render"
|
|
CATEGORY = "🔁 Sequential Batcher/Image"
|
|
|
|
def render(self, progress, padding, width, height, colour, background):
|
|
from PIL import Image, ImageDraw
|
|
img = Image.new("RGB", (width, height), background)
|
|
draw = ImageDraw.Draw(img)
|
|
draw.rectangle((padding, padding, width - padding - 1, height - padding - 1), outline=colour)
|
|
if progress > 0:
|
|
ip = padding + 2
|
|
draw.rectangle((ip, ip, max(ip+1, (width - ip - 1) * progress), height - ip - 1), outline=colour, fill=colour)
|
|
tensor = torch.from_numpy(np.array(img).astype(np.float32) / 255.0).unsqueeze(0)
|
|
return (tensor,)
|
|
|
|
|
|
@register_node
|
|
class ImageBatchToList:
|
|
"""Splits an image batch into a list for iteration."""
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {"required": {"images": ("IMAGE",)}}
|
|
RETURN_TYPES = ("IMAGE",)
|
|
RETURN_NAMES = ("image",)
|
|
OUTPUT_IS_LIST = (True,)
|
|
FUNCTION = "go"
|
|
CATEGORY = "🔁 Sequential Batcher/Image"
|
|
|
|
def go(self, images):
|
|
return ([images[i:i+1] for i in range(images.shape[0])],)
|
|
|
|
|
|
@register_node
|
|
class ImageListToBatch:
|
|
"""Gathers an iterated list of images into a single batch."""
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {"required": {"images": ("IMAGE",)}}
|
|
INPUT_IS_LIST = True
|
|
RETURN_TYPES = ("IMAGE",)
|
|
RETURN_NAMES = ("images",)
|
|
FUNCTION = "go"
|
|
CATEGORY = "🔁 Sequential Batcher/Image"
|
|
|
|
def go(self, images):
|
|
return (torch.cat(images, dim=0),)
|
|
|
|
|
|
@register_node
|
|
class LatentBatchToList:
|
|
"""Splits a latent batch into a list for iteration."""
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {"required": {"latents": ("LATENT",)}}
|
|
RETURN_TYPES = ("LATENT",)
|
|
RETURN_NAMES = ("latent",)
|
|
OUTPUT_IS_LIST = (True,)
|
|
FUNCTION = "go"
|
|
CATEGORY = "🔁 Sequential Batcher/Latent"
|
|
|
|
def go(self, latents):
|
|
samples = latents["samples"]
|
|
return ([{"samples": samples[i:i+1]} for i in range(samples.shape[0])],)
|
|
|
|
|
|
@register_node
|
|
class LatentListToBatch:
|
|
"""Gathers an iterated list of latents into a single batch."""
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {"required": {"latents": ("LATENT",)}}
|
|
INPUT_IS_LIST = True
|
|
RETURN_TYPES = ("LATENT",)
|
|
RETURN_NAMES = ("latents",)
|
|
FUNCTION = "go"
|
|
CATEGORY = "🔁 Sequential Batcher/Latent"
|
|
|
|
def go(self, latents):
|
|
return ({"samples": torch.cat([l["samples"] for l in latents], dim=0)},)
|