diff --git a/nodes/graph_utils.py b/nodes/graph_utils.py index 89341ba..7c35ec9 100644 --- a/nodes/graph_utils.py +++ b/nodes/graph_utils.py @@ -244,4 +244,31 @@ class FitNumber: return (res,) -__nodes__ = [StringReplace, FitNumber, GetBatchFromHistory, AnyToString] +class ConcatImages: + """Add images to batch""" + + RETURN_TYPES = ("IMAGE",) + FUNCTION = "concatenate_tensors" + CATEGORY = "mtb/image" + + @classmethod + def INPUT_TYPES(cls): + return { + "required": {"reverse": ("BOOLEAN", {"default": False})}, + } + + def concatenate_tensors(self, reverse, **kwargs): + tensors = tuple(kwargs.values()) + batch_sizes = [tensor.size(0) for tensor in tensors] + + concatenated = torch.cat(tensors, dim=0) + + # Update the batch size in the concatenated tensor + concatenated_size = list(concatenated.size()) + concatenated_size[0] = sum(batch_sizes) + concatenated = concatenated.view(*concatenated_size) + + return (concatenated,) + + +__nodes__ = [StringReplace, FitNumber, GetBatchFromHistory, AnyToString, ConcatImages] diff --git a/nodes/image_interpolation.py b/nodes/image_interpolation.py index 6bd39ae..0774952 100644 --- a/nodes/image_interpolation.py +++ b/nodes/image_interpolation.py @@ -1,16 +1,18 @@ -from typing import List -from pathlib import Path -import os import glob -import folder_paths -from ..log import log -import torch -from frame_interpolation.eval import util, interpolator -import numpy as np +import os +from pathlib import Path +from typing import List + import comfy -import comfy.utils -import tensorflow as tf import comfy.model_management as model_management +import comfy.utils +import folder_paths +import numpy as np +import tensorflow as tf +import torch +from frame_interpolation.eval import interpolator, util + +from ..log import log class LoadFilmModel: @@ -114,41 +116,4 @@ class FilmInterpolation: return (out_tensors,) -class ConcatImages: - """Add images to batch""" - - RETURN_TYPES = ("IMAGE",) - FUNCTION = "concat_images" - CATEGORY = "mtb/image" - - @classmethod - def INPUT_TYPES(cls): - return { - "required": { - "imageA": ("IMAGE",), - "imageB": ("IMAGE",), - }, - } - - @classmethod - def concatenate_tensors(cls, A: torch.Tensor, B: torch.Tensor): - # Get the batch sizes of A and B - batch_size_A = A.size(0) - batch_size_B = B.size(0) - - # Concatenate the tensors along the batch dimension - concatenated = torch.cat((A, B), dim=0) - - # Update the batch size in the concatenated tensor - concatenated_size = list(concatenated.size()) - concatenated_size[0] = batch_size_A + batch_size_B - concatenated = concatenated.view(*concatenated_size) - - return concatenated - - def concat_images(self, imageA: torch.Tensor, imageB: torch.Tensor): - log.debug(f"Concatenating A ({imageA.shape}) and B ({imageB.shape})") - return (self.concatenate_tensors(imageA, imageB),) - - -__nodes__ = [LoadFilmModel, FilmInterpolation, ConcatImages] +__nodes__ = [LoadFilmModel, FilmInterpolation] diff --git a/web/mtb_widgets.js b/web/mtb_widgets.js index 3f147a2..03918f0 100644 --- a/web/mtb_widgets.js +++ b/web/mtb_widgets.js @@ -877,7 +877,8 @@ const mtb_widgets = { break } - case 'Stack Images (mtb)': { + case 'Stack Images (mtb)': + case 'Concat Images (mtb)': { shared.setupDynamicConnections(nodeType, 'image', 'IMAGE') break