diff --git a/__init__.py b/__init__.py index 8b77090..bc0556a 100644 --- a/__init__.py +++ b/__init__.py @@ -9,20 +9,25 @@ NODE_CLASS_MAPPINGS = { "FS: Fit Image And Resize": FitResizeLatent, "FS: Load Image And Resize To Fit": LoadToFitResizeLatent, "FS: Pick Image From Batch": RandomImageFromBatch, + "FS: Pick Image From Batches": RandomImageFromBatches, + "FS: Pick Image From List": RandomImageFromList, "FS: Crop Image Into Even Pieces": CropImageIntoEvenPieces, + "FS: Image Region To Mask": ImageRegionMask, } NODE_DISPLAY_NAME_MAPPINGS = { - "FS: Fit Size From Int": "Fit Size From Int", - "FS: Fit Size From Image": "Fit Size From Image", - "FS: Fit Image And Resize": "Fit Image And Resize", - "FS: Load Image And Resize To Fit": "Load Image And Resize To Fit", - "FS: Pick Image From Batch": "Pick Image From Batch", - "FS: Crop Image Into Even Pieces": "Crop Image Into Even Pieces", + "FS: Fit Size From Int": "Fit Size From Int (FS)", + "FS: Fit Size From Image": "Fit Size From Image (FS)", + "FS: Fit Image And Resize": "Fit Image And Resize (FS)", + "FS: Load Image And Resize To Fit": "Load Image And Resize To Fit (FS)", + "FS: Pick Image From Batch": "Pick Image From Batch (FS)", + "FS: Pick Image From Batches": "Pick Image From Batches (FS)", + "FS: Pick Image From List": "Pick Image From List (FS)", + "FS: Crop Image Into Even Pieces": "Crop Image Into Even Pieces (FS)", + "FS: Image Region To Mask": "Image Region To Mask (FS)", } __all__ = ['NODE_CLASS_MAPPINGS', 'NODE_DISPLAY_NAME_MAPPINGS'] EXTENSION_NAME = "Fitsize" - -symlink_web_dir("js", EXTENSION_NAME) \ No newline at end of file +WEB_DIRECTORY = "./js" diff --git a/js/widgets.js b/js/widgets.js index e69de29..1396bfa 100644 --- a/js/widgets.js +++ b/js/widgets.js @@ -0,0 +1,87 @@ +import { app } from '../../scripts/app.js' + +// from: https://github.com/melMass/comfy_mtb + +export const setupDynamicConnections = (nodeType, prefix, inputType) => { + const onNodeCreated = nodeType.prototype.onNodeCreated + nodeType.prototype.onNodeCreated = function () { + const r = onNodeCreated ? onNodeCreated.apply(this, arguments) : undefined + console.log('onNodeCreated', `${prefix}_1`) + this.addInput(`${prefix}_1`, inputType) + return r + } + + const onConnectionsChange = nodeType.prototype.onConnectionsChange + nodeType.prototype.onConnectionsChange = function ( + type, + index, + connected, + link_info + ) { + const r = onConnectionsChange + ? onConnectionsChange.apply(this, arguments) + : undefined + dynamic_connection(this, index, connected, `${prefix}_`, inputType) + } +} + +export const dynamic_connection = ( + node, + index, + connected, + connectionPrefix = 'input_', + connectionType = 'PSDLAYER', + nameArray = [] +) => { + if (!node.inputs[index].name.startsWith(connectionPrefix)) { + return + } + // remove all non connected inputs + if (!connected && node.inputs.length > 1) { + if (node.widgets) { + const w = node.widgets.find((w) => w.name === node.inputs[index].name) + if (w) { + w.onRemoved?.() + node.widgets.length = node.widgets.length - 1 + } + } + node.removeInput(index) + + // make inputs sequential again + for (let i = 0; i < node.inputs.length; i++) { + const name = + i < nameArray.length ? nameArray[i] : `${connectionPrefix}${i + 1}` + node.inputs[i].label = name + node.inputs[i].name = name + } + } + + // add an extra input + if (node.inputs[node.inputs.length - 1].link != undefined) { + const nextIndex = node.inputs.length + const name = + nextIndex < nameArray.length + ? nameArray[nextIndex] + : `${connectionPrefix}${nextIndex + 1}` + + node.addInput(name, connectionType) + } +} + + + + + + + +app.registerExtension({ + name: "Comfy.Fitsize.PickImageFromBatches", + async beforeRegisterNodeDef(nodeType, nodeData, app) { + if (!nodeData.name.startsWith('FS:')) { + return + } + if(nodeData.name === "FS: Pick Image From Batches") { + setupDynamicConnections(nodeType, 'batches', 'IMAGE') + } + } +}) \ No newline at end of file diff --git a/nodes.py b/nodes.py index c0d4d16..c368e76 100644 --- a/nodes.py +++ b/nodes.py @@ -1,3 +1,4 @@ +import random from PIL import Image, ImageOps import torch import os @@ -345,10 +346,55 @@ class CropImageIntoEvenPieces: crop = image[: , y : y + crop_height , x : x + crop_width , :] pieces.append(torch.from_numpy(crop)) - # image[:, y : y + height, x : x + width, :] return (torch.cat(pieces, dim=0), ) +class ImageRegionMask: + def __init__(self): + pass + + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "image": ("IMAGE",), + "rows": ("INT", {"default": 3, "min": 1, "max": 32, "step": 1,}), + "columns": ("INT", {"default": 1, "min": 1, "max": 32, "step": 1,}), + "chosen_row": ("INT", {"default": 0, "min": 0, "max": 32, "step": 1,}), + "chosen_column": ("INT", {"default": 0, "min": 0, "max": 32, "step": 1,}), + }, + } + + RETURN_TYPES = ("MASK",) + + FUNCTION = "run" + + CATEGORY = "Fitsize/Mask" + + def run(self, image, rows, columns, chosen_row, chosen_column): + + if rows < 1: + rows = 1 + if columns < 1: + columns = 1 + + w = image.shape[2] # width + h = image.shape[1] # height + + crop_width = int(w / columns) + crop_height = int(h / rows) + + mask = torch.zeros((h, w)) + + min_y = crop_height * chosen_row + max_y = min_y + crop_height + min_x = crop_width * chosen_column + max_x = min_x + crop_width + + mask[int(min_y):int(max_y), int(min_x):int(max_x)] = 1 + + return (mask.unsqueeze(0), ) + class RandomImageFromBatch: @@ -387,15 +433,107 @@ class RandomImageFromBatch: if select_amount < 1: select_amount = 1 - - print( - f"RandomImageFromBatch: start_index {start_index},", - f"select_amount {select_amount}", - f"total_images {images.shape[0]}",) - selected = images[start_index:start_index + select_amount] - print(f"RandomImageFromBatch: selected {selected.shape[0]} images") + return (selected, ) + +class RandomImageFromList: + def __init__(self): + pass + + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "list": ("IMAGE", ), + "seed": ("INT", {"default": 0}), + "start_index": ("INT", {"default": -1, "min": -1, "max": 32, "step": 1,}), + "select_amount": ("INT", {"default": 1, "min": 1, "max": 32, "step": 1,}), + }, + } + + RETURN_TYPES = ("IMAGE",) + + FUNCTION = "run" + + CATEGORY = "Fitsize/Image" + + def run(self, list, seed, start_index, select_amount): + print(f'type of list: {type(list)}, length: {len(list)}') + + list_length = len(list) + + if start_index == -1: + start_index = np.random.randint(0, list_length) + # return random.choice(list, select_amount) + if start_index >= list_length: + start_index = list_length-1 + + if select_amount > list_length: + select_amount = list_length + if select_amount < 1: + select_amount = 1 + + selected = list[start_index:start_index + select_amount] + + print(f'selected: {start_index} to {start_index + select_amount} found {len(selected)}') return (selected, ) + + +class RandomImageFromBatches: + def __init__(self): + pass + + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "seed": ("INT",{"default": 0}), + "start_index": ("INT", {"default": -1, "min": -1, "max": 32, "step": 1,}), + "select_amount": ("INT", {"default": 1, "min": 1, "max": 32, "step": 1,}), + }, + } + + RETURN_TYPES = ("IMAGE",) + + FUNCTION = "run" + + CATEGORY = "Fitsize/Image" + + def run(self, seed=0, start_index=0, select_amount=1, **kwargs): + + batches = kwargs.values() + + selected = [] + + print(f'len(batches): {len(batches)}') + + + + for img in batches: + + # if type(images) == torch.Tensor: + # images = images.numpy() + + if start_index == -1: + start_index = np.random.randint(0, img.shape[0]) + if start_index >= img.shape[0]: + start_index = img.shape[0]-1 + + if select_amount > img.shape[0]: + select_amount = img.shape[0] + if select_amount < 1: + select_amount = 1 + + # add images to selected + selected.append(img[start_index:start_index + select_amount]) + + # try to return a tensor of images if all widths and heights match + try: + selected = torch.cat(selected, dim=0) + except: + pass + + return (selected, )