diff --git a/essentials.py b/essentials.py index 45219c0..f2e5f82 100644 --- a/essentials.py +++ b/essentials.py @@ -226,7 +226,7 @@ class ImageEnhanceDifference: "required": { "image1": ("IMAGE",), "image2": ("IMAGE",), - "exponent": ("FLOAT", { "default": 0.50, "min": 0.00, "max": 1.00, "step": 0.05, "display": "number" }), + "exponent": ("FLOAT", { "default": 0.75, "min": 0.00, "max": 1.00, "step": 0.05, "display": "number" }), } } @@ -263,10 +263,10 @@ class MaskFlip: def execute(self, mask, axis): dim = () if "y" in axis: - dim += (0,) - if "x" in axis: dim += (1,) - mask = torch.flip(mask, dim) + if "x" in axis: + dim += (2,) + mask = torch.flip(mask, dims=dim) return(mask,) @@ -320,6 +320,133 @@ class MaskPreview(SaveImage): return( results ) +class MaskBatch: + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "mask1": ("MASK",), + "mask2": ("MASK",), + } + } + + RETURN_TYPES = ("MASK",) + FUNCTION = "execute" + CATEGORY = "essentials" + + def execute(self, mask1, mask2): + if mask1.shape[1:] != mask2.shape[1:]: + mask2 = F.interpolate(mask2.unsqueeze(1), size=(mask1.shape[1], mask1.shape[2]), mode="bicubic").squeeze(1) + + out = torch.cat((mask1, mask2), dim=0) + return (out,) + +def cubic_bezier(t, p): + p0, p1, p2, p3 = p + return (1 - t)**3 * p0 + 3 * (1 - t)**2 * t * p1 + 3 * (1 - t) * t**2 * p2 + t**3 * p3 + +class TransitionMask: + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "width": ("INT", { "default": 512, "min": 1, "max": MAX_RESOLUTION, "step": 1, "display": "number" }), + "height": ("INT", { "default": 512, "min": 1, "max": MAX_RESOLUTION, "step": 1, "display": "number" }), + "frames": ("INT", { "default": 16, "min": 1, "max": 9999, "step": 1, "display": "number" }), + "start_frame": ("INT", { "default": 0, "min": 0, "step": 1, "display": "number" }), + "end_frame": ("INT", { "default": 9999, "min": 0, "step": 1, "display": "number" }), + "transition_type": (["horizontal slide", "vertical slide", "horizontal bar", "vertical bar", "center box", "horizontal door", "vertical door", "circle", "fade"],), + "timing_function": (["linear", "in", "out", "in-out"],) + } + } + + RETURN_TYPES = ("MASK",) + FUNCTION = "execute" + CATEGORY = "essentials" + + def execute(self, width, height, frames, start_frame, end_frame, transition_type, timing_function): + if timing_function == 'in': + tf = [0.0, 0.0, 0.5, 1.0] + elif timing_function == 'out': + tf = [0.0, 0.5, 1.0, 1.0] + elif timing_function == 'in-out': + tf = [0, 1, 0, 1] + #elif timing_function == 'back': + # tf = [0, 1.334, 1.334, 0] + else: + tf = [0, 0, 1, 1] + + out = [] + + end_frame = min(frames, end_frame) + transition = end_frame - start_frame + + if start_frame > 0: + out = out + [torch.full((height, width), 0.0, dtype=torch.float32, device="cpu")] * start_frame + + for i in range(transition): + frame = torch.full((height, width), 0.0, dtype=torch.float32, device="cpu") + progress = i/(transition-1) + + if timing_function != 'linear': + progress = cubic_bezier(progress, tf) + + if "horizontal slide" in transition_type: + pos = round(width*progress) + frame[:, :pos] = 1.0 + elif "vertical slide" in transition_type: + pos = round(height*progress) + frame[:pos, :] = 1.0 + elif "box" in transition_type: + box_w = round(width*progress) + box_h = round(height*progress) + x1 = (width - box_w) // 2 + y1 = (height - box_h) // 2 + x2 = x1 + box_w + y2 = y1 + box_h + frame[y1:y2, x1:x2] = 1.0 + elif "circle" in transition_type: + radius = math.ceil(math.sqrt(pow(width,2)+pow(height,2))*progress/2) + c_x = width // 2 + c_y = height // 2 + # is this real life? Am I hallucinating? + x = torch.arange(0, width, dtype=torch.float32, device="cpu") + y = torch.arange(0, height, dtype=torch.float32, device="cpu") + y, x = torch.meshgrid((y, x), indexing="ij") + circle = ((x - c_x) ** 2 + (y - c_y) ** 2) <= (radius ** 2) + frame[circle] = 1.0 + elif "horizontal bar" in transition_type: + bar = round(height*progress) + y1 = (height - bar) // 2 + y2 = y1 + bar + frame[y1:y2, :] = 1.0 + elif "vertical bar" in transition_type: + bar = round(width*progress) + x1 = (width - bar) // 2 + x2 = x1 + bar + frame[:, x1:x2] = 1.0 + elif "horizontal door" in transition_type: + bar = math.ceil(height*progress/2) + if bar > 0: + frame[:bar, :] = 1.0 + frame[-bar:, :] = 1.0 + elif "vertical door" in transition_type: + bar = math.ceil(width*progress/2) + if bar > 0: + frame[:, :bar] = 1.0 + frame[:, -bar:] = 1.0 + elif "fade" in transition_type: + frame[:,:] = progress + + out.append(frame) + + if end_frame < frames: + out = out + [torch.full((height, width), 1.0, dtype=torch.float32, device="cpu")] * (frames - end_frame) + + out = torch.stack(out, dim=0) + + return (out, ) + def min_(tensor_list): # return the element-wise min of the tensor list. x = torch.stack(tensor_list) @@ -494,6 +621,8 @@ NODE_CLASS_MAPPINGS = { "MaskBlur+": MaskBlur, "MaskFlip+": MaskFlip, "MaskPreview+": MaskPreview, + "MaskBatch+": MaskBatch, + "TransitionMask+": TransitionMask, "SimpleMath+": SimpleMath, "ConsoleDebug+": ConsoleDebug, @@ -516,6 +645,8 @@ NODE_DISPLAY_NAME_MAPPINGS = { "MaskBlur+": "🔧 Mask Blur", "MaskFlip+": "🔧 Mask Flip", "MaskPreview+": "🔧 Mask Preview", + "MaskBatch+": "🔧 Mask Batch", + "TransitionMask+": "🔧 Transition Mask", "SimpleMath+": "🔧 Simple Math", "ConsoleDebug+": "🔧 Console Debug",