custom_nodes/comfy_mtb/nodes/batch.py

from io import BytesIO

import cv2
import numpy as np
import torch
from PIL import Image

from ..log import log
from ..utils import EASINGS, apply_easing, pil2tensor
from .transform import MTB_TransformImage


def hex_to_rgb(hex_color: str, bgr: bool = False):
    hex_color = hex_color.lstrip("#")
    if bgr:
        return tuple(int(hex_color[i : i + 2], 16) for i in (4, 2, 0))

    return tuple(int(hex_color[i : i + 2], 16) for i in (0, 2, 4))


class MTB_BatchFloatMath:
    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "reverse": ("BOOLEAN", {"default": False}),
                "operation": (
                    ["add", "sub", "mul", "div", "pow", "abs"],
                    {"default": "add"},
                ),
            }
        }

    RETURN_TYPES = ("FLOATS",)
    CATEGORY = "mtb/utils"
    FUNCTION = "execute"

    def execute(self, reverse: bool, operation: str, **kwargs: list[float]):
        res: list[float] = []
        vals = list(kwargs.values())

        if reverse:
            vals = vals[::-1]

        ref_count = len(vals[0])
        for v in vals:
            if len(v) != ref_count:
                raise ValueError(
                    f"All values must have the same length (current: {len(v)}, ref: {ref_count}"
                )

        match operation:
            case "add":
                for i in range(ref_count):
                    result = sum(v[i] for v in vals)
                    res.append(result)
            case "sub":
                for i in range(ref_count):
                    result = vals[0][i] - sum(v[i] for v in vals[1:])
                    res.append(result)
            case "mul":
                for i in range(ref_count):
                    result = vals[0][i] * vals[1][i]
                    res.append(result)
            case "div":
                for i in range(ref_count):
                    result = vals[0][i] / vals[1][i]
                    res.append(result)
            case "pow":
                for i in range(ref_count):
                    result: float = vals[0][i] ** vals[1][i]
                    res.append(result)
            case "abs":
                for i in range(ref_count):
                    result = abs(vals[0][i])
                    res.append(result)
            case _:
                log.info(f"For now this mode ({operation}) is not implemented")

        return (res,)


class MTB_BatchFloatNormalize:
    """Normalize the values in the list of floats"""

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {"floats": ("FLOATS",)},
        }

    RETURN_TYPES = ("FLOATS",)
    RETURN_NAMES = ("normalized_floats",)
    CATEGORY = "mtb/batch"
    FUNCTION = "execute"

    def execute(
        self,
        floats: list[float],
    ):
        min_value = min(floats)
        max_value = max(floats)

        normalized_floats = [
            (x - min_value) / (max_value - min_value) for x in floats
        ]
        log.debug(f"Floats: {floats}")
        log.debug(f"Normalized Floats: {normalized_floats}")

        return (normalized_floats,)


class MTB_BatchTimeWrap:
    """Remap a batch using a time curve (FLOATS)"""

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "target_count": ("INT", {"default": 25, "min": 2}),
                "frames": ("IMAGE",),
                "curve": ("FLOATS",),
            },
        }

    RETURN_TYPES = ("IMAGE", "FLOATS")
    RETURN_NAMES = ("image", "interpolated_floats")
    CATEGORY = "mtb/batch"
    FUNCTION = "execute"

    def execute(
        self, target_count: int, frames: torch.Tensor, curve: list[float]
    ):
        """Apply time warping to a list of video frames based on a curve."""
        log.debug(f"Input frames shape: {frames.shape}")
        log.debug(f"Curve: {curve}")

        total_duration = sum(curve)

        log.debug(f"Total duration: {total_duration}")

        B, H, W, C = frames.shape

        log.debug(f"Batch Size: {B}")

        normalized_times = np.linspace(0, 1, target_count)
        interpolated_curve = np.interp(
            normalized_times, np.linspace(0, 1, len(curve)), curve
        ).tolist()
        log.debug(f"Interpolated curve: {interpolated_curve}")

        interpolated_frame_indices = [
            (B - 1) * value for value in interpolated_curve
        ]
        log.debug(f"Interpolated frame indices: {interpolated_frame_indices}")

        rounded_indices = [
            int(round(idx)) for idx in interpolated_frame_indices
        ]
        rounded_indices = np.clip(rounded_indices, 0, B - 1)

        # Gather frames based on interpolated indices
        warped_frames = []
        for index in rounded_indices:
            warped_frames.append(frames[index].unsqueeze(0))

        warped_tensor = torch.cat(warped_frames, dim=0)
        log.debug(f"Warped frames shape: {warped_tensor.shape}")
        return (warped_tensor, interpolated_curve)


class MTB_BatchMake:
    """Simply duplicates the input frame as a batch"""

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "image": ("IMAGE",),
                "count": ("INT", {"default": 1}),
            }
        }

    RETURN_TYPES = ("IMAGE",)
    FUNCTION = "generate_batch"
    CATEGORY = "mtb/batch"

    def generate_batch(self, image: torch.Tensor, count):
        if len(image.shape) == 3:
            image = image.unsqueeze(0)

        return (image.repeat(count, 1, 1, 1),)


class MTB_BatchShape:
    """Generates a batch of 2D shapes with optional shading (experimental)"""

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "count": ("INT", {"default": 1}),
                "shape": (
                    ["Box", "Circle", "Diamond", "Tube"],
                    {"default": "Circle"},
                ),
                "image_width": ("INT", {"default": 512}),
                "image_height": ("INT", {"default": 512}),
                "shape_size": ("INT", {"default": 100}),
                "color": ("COLOR", {"default": "#ffffff"}),
                "bg_color": ("COLOR", {"default": "#000000"}),
                "shade_color": ("COLOR", {"default": "#000000"}),
                "thickness": ("INT", {"default": 5}),
                "shadex": ("FLOAT", {"default": 0.0}),
                "shadey": ("FLOAT", {"default": 0.0}),
            },
        }

    RETURN_TYPES = ("IMAGE",)
    FUNCTION = "generate_shapes"
    CATEGORY = "mtb/batch"

    def generate_shapes(
        self,
        count,
        shape,
        image_width,
        image_height,
        shape_size,
        color,
        bg_color,
        shade_color,
        thickness,
        shadex,
        shadey,
    ):
        log.debug(f"COLOR: {color}")
        log.debug(f"BG_COLOR: {bg_color}")
        log.debug(f"SHADE_COLOR: {shade_color}")

        # Parse color input to BGR tuple for OpenCV
        color = hex_to_rgb(color)
        bg_color = hex_to_rgb(bg_color)
        shade_color = hex_to_rgb(shade_color)
        res = []
        for x in range(count):
            # Initialize an image canvas
            canvas = np.full(
                (image_height, image_width, 3), bg_color, dtype=np.uint8
            )
            mask = np.zeros((image_height, image_width), dtype=np.uint8)

            # Compute the center point of the shape
            center = (image_width // 2, image_height // 2)

            if shape == "Box":
                half_size = shape_size // 2
                top_left = (center[0] - half_size, center[1] - half_size)
                bottom_right = (center[0] + half_size, center[1] + half_size)
                cv2.rectangle(mask, top_left, bottom_right, 255, -1)
            elif shape == "Circle":
                cv2.circle(mask, center, shape_size // 2, 255, -1)
            elif shape == "Diamond":
                pts = np.array(
                    [
                        [center[0], center[1] - shape_size // 2],
                        [center[0] + shape_size // 2, center[1]],
                        [center[0], center[1] + shape_size // 2],
                        [center[0] - shape_size // 2, center[1]],
                    ]
                )
                cv2.fillPoly(mask, [pts], 255)

            elif shape == "Tube":
                cv2.ellipse(
                    mask,
                    center,
                    (shape_size // 2, shape_size // 2),
                    0,
                    0,
                    360,
                    255,
                    thickness,
                )

            # Color the shape
            canvas[mask == 255] = color

            # Apply shading effects to a separate shading canvas
            shading = np.zeros_like(canvas, dtype=np.float32)
            shading[:, :, 0] = shadex * np.linspace(0, 1, image_width)
            shading[:, :, 1] = shadey * np.linspace(
                0, 1, image_height
            ).reshape(-1, 1)
            shading_canvas = cv2.addWeighted(
                canvas.astype(np.float32), 1, shading, 1, 0
            ).astype(np.uint8)

            # Apply shading only to the shape area using the mask
            canvas[mask == 255] = shading_canvas[mask == 255]
            res.append(canvas)

        return (pil2tensor(res),)


class MTB_BatchFloatFill:
    """Fills a batch float with a single value until it reaches the target length"""

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "floats": ("FLOATS",),
                "direction": (["head", "tail"], {"default": "tail"}),
                "value": ("FLOAT", {"default": 0.0}),
                "count": ("INT", {"default": 1}),
            }
        }

    FUNCTION = "fill_floats"
    RETURN_TYPES = ("FLOATS",)
    CATEGORY = "mtb/batch"

    def fill_floats(self, floats, direction, value, count):
        size = len(floats)
        if size > count:
            raise ValueError(
                f"Size ({size}) is less then target count ({count})"
            )

        rem = count - size
        if direction == "tail":
            floats = floats + [value] * rem
        else:
            floats = [value] * rem + floats
        return (floats,)


class MTB_BatchFloatAssemble:
    """Assembles mutiple batches of floats into a single stream (batch)"""

    @classmethod
    def INPUT_TYPES(cls):
        return {"required": {"reverse": ("BOOLEAN", {"default": False})}}

    RETURN_TYPES = ("FLOATS",)
    CATEGORY = "mtb/batch"
    FUNCTION = "assemble_floats"

    def assemble_floats(self, reverse: bool, **kwargs: list[float]):
        res: list[float] = []

        if reverse:
            for x in reversed(kwargs.values()):
                if x:
                    res += x
        else:
            for x in kwargs.values():
                if x:
                    res += x

        return (res,)


class MTB_BatchFloat:
    """Generates a batch of float values with interpolation"""

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "mode": (
                    ["Single", "Steps"],
                    {"default": "Steps"},
                ),
                "count": ("INT", {"default": 2}),
                "min": ("FLOAT", {"default": 0.0, "step": 0.001}),
                "max": ("FLOAT", {"default": 1.0, "step": 0.001}),
                "easing": (
                    [
                        "Linear",
                        "Sine In",
                        "Sine Out",
                        "Sine In/Out",
                        "Quart In",
                        "Quart Out",
                        "Quart In/Out",
                        "Cubic In",
                        "Cubic Out",
                        "Cubic In/Out",
                        "Circ In",
                        "Circ Out",
                        "Circ In/Out",
                        "Back In",
                        "Back Out",
                        "Back In/Out",
                        "Elastic In",
                        "Elastic Out",
                        "Elastic In/Out",
                        "Bounce In",
                        "Bounce Out",
                        "Bounce In/Out",
                    ],
                    {"default": "Linear"},
                ),
            }
        }

    FUNCTION = "set_floats"
    RETURN_TYPES = ("FLOATS",)
    CATEGORY = "mtb/batch"

    def set_floats(self, mode, count, min, max, easing):
        if mode == "Steps" and count == 1:
            raise ValueError(
                "Steps mode requires at least a count of 2 values"
            )
        keyframes = []
        if mode == "Single":
            keyframes = [min] * count
            return (keyframes,)

        for i in range(count):
            normalized_step = i / (count - 1)
            eased_step = apply_easing(normalized_step, easing)
            eased_value = min + (max - min) * eased_step
            keyframes.append(eased_value)

        return (keyframes,)


class MTB_BatchMerge:
    """Merges multiple image batches with different frame counts"""

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "fusion_mode": (
                    ["add", "multiply", "average"],
                    {"default": "average"},
                ),
                "fill": (["head", "tail"], {"default": "tail"}),
            }
        }

    RETURN_TYPES = ("IMAGE",)
    FUNCTION = "merge_batches"
    CATEGORY = "mtb/batch"

    def merge_batches(self, fusion_mode: str, fill: str, **kwargs):
        images = kwargs.values()
        max_frames = max(img.shape[0] for img in images)

        adjusted_images = []
        for img in images:
            frame_count = img.shape[0]
            if frame_count < max_frames:
                fill_frame = img[0] if fill == "head" else img[-1]
                fill_frames = fill_frame.repeat(
                    max_frames - frame_count, 1, 1, 1
                )
                adjusted_batch = (
                    torch.cat((fill_frames, img), dim=0)
                    if fill == "head"
                    else torch.cat((img, fill_frames), dim=0)
                )
            else:
                adjusted_batch = img
            adjusted_images.append(adjusted_batch)

        # Merge the adjusted batches
        merged_image = None
        for img in adjusted_images:
            if merged_image is None:
                merged_image = img
            else:
                if fusion_mode == "add":
                    merged_image += img
                elif fusion_mode == "multiply":
                    merged_image *= img
                elif fusion_mode == "average":
                    merged_image = (merged_image + img) / 2

        return (merged_image,)


class MTB_Batch2dTransform:
    """Transform a batch of images using a batch of keyframes"""

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "image": ("IMAGE",),
                "border_handling": (
                    ["edge", "constant", "reflect", "symmetric"],
                    {"default": "edge"},
                ),
                "constant_color": ("COLOR", {"default": "#000000"}),
            },
            "optional": {
                "x": ("FLOATS",),
                "y": ("FLOATS",),
                "zoom": ("FLOATS",),
                "angle": ("FLOATS",),
                "shear": ("FLOATS",),
            },
        }

    RETURN_TYPES = ("IMAGE",)
    FUNCTION = "transform_batch"
    CATEGORY = "mtb/batch"

    def get_num_elements(
        self, param: None | torch.Tensor | list[torch.Tensor] | list[float]
    ) -> int:
        if isinstance(param, torch.Tensor):
            return torch.numel(param)

        elif isinstance(param, list):
            return len(param)

        return 0

    def transform_batch(
        self,
        image: torch.Tensor,
        border_handling: str,
        constant_color: str,
        x: list[float] | None = None,
        y: list[float] | None = None,
        zoom: list[float] | None = None,
        angle: list[float] | None = None,
        shear: list[float] | None = None,
    ):
        if all(
            self.get_num_elements(param) <= 0
            for param in [x, y, zoom, angle, shear]
        ):
            raise ValueError(
                "At least one transform parameter must be provided"
            )

        keyframes: dict[str, list[float]] = {
            "x": [],
            "y": [],
            "zoom": [],
            "angle": [],
            "shear": [],
        }

        default_vals = {"x": 0, "y": 0, "zoom": 1.0, "angle": 0, "shear": 0}

        if x and self.get_num_elements(x) > 0:
            keyframes["x"] = x
        if y and self.get_num_elements(y) > 0:
            keyframes["y"] = y
        if zoom and self.get_num_elements(zoom) > 0:
            # some easing types like elastic can pull back... maybe it should abs the value?
            keyframes["zoom"] = [max(x, 0.00001) for x in zoom]
        if angle and self.get_num_elements(angle) > 0:
            keyframes["angle"] = angle
        if shear and self.get_num_elements(shear) > 0:
            keyframes["shear"] = shear

        for name, values in keyframes.items():
            count = len(values)
            if count > 0 and count != image.shape[0]:
                raise ValueError(
                    f"Length of {name} values ({count}) must match number of images ({image.shape[0]})"
                )
            if count == 0:
                keyframes[name] = [default_vals[name]] * image.shape[0]

        transformer = MTB_TransformImage()
        res = [
            transformer.transform(
                image[i].unsqueeze(0),
                keyframes["x"][i],
                keyframes["y"][i],
                keyframes["zoom"][i],
                keyframes["angle"][i],
                keyframes["shear"][i],
                border_handling,
                constant_color,
            )[0]
            for i in range(image.shape[0])
        ]
        return (torch.cat(res, dim=0),)


class MTB_BatchFloatFit:
    """Fit a list of floats using a source and target range"""

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "values": ("FLOATS", {"forceInput": True}),
                "clamp": ("BOOLEAN", {"default": False}),
                "auto_compute_source": ("BOOLEAN", {"default": False}),
                "source_min": ("FLOAT", {"default": 0.0, "step": 0.01}),
                "source_max": ("FLOAT", {"default": 1.0, "step": 0.01}),
                "target_min": ("FLOAT", {"default": 0.0, "step": 0.01}),
                "target_max": ("FLOAT", {"default": 1.0, "step": 0.01}),
                "easing": (
                    EASINGS,
                    {"default": "Linear"},
                ),
            }
        }

    FUNCTION = "fit_range"
    RETURN_TYPES = ("FLOATS",)
    CATEGORY = "mtb/batch"
    DESCRIPTION = "Fit a list of floats using a source and target range"

    def fit_range(
        self,
        values: list[float],
        clamp: bool,
        auto_compute_source: bool,
        source_min: float,
        source_max: float,
        target_min: float,
        target_max: float,
        easing: str,
    ):
        if auto_compute_source:
            source_min = min(values)
            source_max = max(values)

        from .graph_utils import MTB_FitNumber

        res = []
        fit_number = MTB_FitNumber()
        for value in values:
            (transformed_value,) = fit_number.set_range(
                value,
                clamp,
                source_min,
                source_max,
                target_min,
                target_max,
                easing,
            )
            res.append(transformed_value)

        return (res,)


class MTB_PlotBatchFloat:
    """Plot floats"""

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "width": ("INT", {"default": 768}),
                "height": ("INT", {"default": 768}),
                "point_size": ("INT", {"default": 4}),
                "seed": ("INT", {"default": 1}),
                "start_at_zero": ("BOOLEAN", {"default": False}),
            }
        }

    RETURN_TYPES = ("IMAGE",)
    RETURN_NAMES = ("plot",)
    FUNCTION = "plot"
    CATEGORY = "mtb/batch"

    def plot(
        self,
        width: int,
        height: int,
        point_size: int,
        seed: int,
        start_at_zero: bool,
        interactive_backend: bool = False,
        **kwargs,
    ):
        import matplotlib

        # NOTE: This is for notebook usage or tests, i.e not exposed to comfy that should always use Agg
        if not interactive_backend:
            matplotlib.use("Agg")
        import matplotlib.pyplot as plt

        fig, ax = plt.subplots(figsize=(width / 100, height / 100), dpi=100)
        fig.set_edgecolor("black")
        fig.patch.set_facecolor("#2e2e2e")
        # Setting background color and grid
        ax.set_facecolor("#2e2e2e")  # Dark gray background
        ax.grid(color="gray", linestyle="-", linewidth=0.5, alpha=0.5)

        # Finding global min and max across all lists for scaling the plot
        all_values = [value for values in kwargs.values() for value in values]
        global_min = min(all_values)
        global_max = max(all_values)

        y_padding = 0.05 * (global_max - global_min)
        ax.set_ylim(global_min - y_padding, global_max + y_padding)

        max_length = max(len(values) for values in kwargs.values())
        if start_at_zero:
            x_values = np.linspace(0, max_length - 1, max_length)
        else:
            x_values = np.linspace(1, max_length, max_length)

        ax.set_xlim(1, max_length)  # Set X-axis limits
        np.random.seed(seed)
        colors = np.random.rand(len(kwargs), 3)  # Generate random RGB values
        for color, (label, values) in zip(colors, kwargs.items()):
            ax.plot(x_values[: len(values)], values, label=label, color=color)
        ax.legend(
            title="Legend",
            title_fontsize="large",
            fontsize="medium",
            edgecolor="black",
            loc="best",
        )

        # Setting labels and title
        ax.set_xlabel("Time", fontsize="large", color="white")
        ax.set_ylabel("Value", fontsize="large", color="white")
        ax.set_title(
            "Plot of Values over Time", fontsize="x-large", color="white"
        )

        # Adjusting tick colors to be visible on dark background
        ax.tick_params(colors="white")

        # Changing color of the axes border
        for _, spine in ax.spines.items():
            spine.set_edgecolor("white")

        # Rendering the plot into a NumPy array
        buf = BytesIO()
        plt.savefig(buf, format="png", bbox_inches="tight")
        buf.seek(0)
        image = Image.open(buf)
        plt.close(fig)  # Closing the figure to free up memory

        return (pil2tensor(image),)

    def draw_point(self, image, point, color, point_size):
        x, y = point
        y = image.shape[0] - 1 - y  # Invert Y-coordinate
        half_size = point_size // 2
        x_start, x_end = (
            max(0, x - half_size),
            min(image.shape[1], x + half_size + 1),
        )
        y_start, y_end = (
            max(0, y - half_size),
            min(image.shape[0], y + half_size + 1),
        )
        image[y_start:y_end, x_start:x_end] = color

    def draw_line(self, image, start, end, color):
        x1, y1 = start
        x2, y2 = end

        # Invert Y-coordinate
        y1 = image.shape[0] - 1 - y1
        y2 = image.shape[0] - 1 - y2

        dx = x2 - x1
        dy = y2 - y1
        is_steep = abs(dy) > abs(dx)
        if is_steep:
            x1, y1 = y1, x1
            x2, y2 = y2, x2
        swapped = False
        if x1 > x2:
            x1, x2 = x2, x1
            y1, y2 = y2, y1
            swapped = True
        dx = x2 - x1
        dy = y2 - y1
        error = int(dx / 2.0)
        y = y1
        ystep = None
        if y1 < y2:
            ystep = 1
        else:
            ystep = -1
        for x in range(x1, x2 + 1):
            coord = (y, x) if is_steep else (x, y)
            image[coord] = color
            error -= abs(dy)
            if error < 0:
                y += ystep
                error += dx
        if swapped:
            image[(x1, y1)] = color
            image[(x2, y2)] = color


DEFAULT_INTERPOLANT = lambda t: t * t * t * (t * (t * 6 - 15) + 10)


class MTB_BatchShake:
    """Applies a shaking effect to batches of images."""

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "images": ("IMAGE",),
                "position_amount_x": ("FLOAT", {"default": 1.0}),
                "position_amount_y": ("FLOAT", {"default": 1.0}),
                "rotation_amount": ("FLOAT", {"default": 10.0}),
                "frequency": ("FLOAT", {"default": 1.0, "min": 0.005}),
                "frequency_divider": ("FLOAT", {"default": 1.0, "min": 0.005}),
                "octaves": ("INT", {"default": 1, "min": 1}),
                "seed": ("INT", {"default": 0}),
            },
        }

    RETURN_TYPES = ("IMAGE", "FLOATS", "FLOATS", "FLOATS")
    RETURN_NAMES = ("image", "pos_x", "pos_y", "rot")
    FUNCTION = "apply_shake"
    CATEGORY = "mtb/batch"

    # def interpolant(self, t):
    # return t * t * t * (t * (t * 6 - 15) + 10)

    def generate_perlin_noise_2d(
        self, shape, res, tileable=(False, False), interpolant=None
    ):
        """Generate a 2D numpy array of perlin noise.

        Args:
            shape: The shape of the generated array (tuple of two ints).
                This must be a multple of res.
            res: The number of periods of noise to generate along each
                axis (tuple of two ints). Note shape must be a multiple of
                res.
            tileable: If the noise should be tileable along each axis
                (tuple of two bools). Defaults to (False, False).
            interpolant: The interpolation function, defaults to
                t*t*t*(t*(t*6 - 15) + 10).

        Returns
        -------
            A numpy array of shape shape with the generated noise.

        Raises
        ------
            ValueError: If shape is not a multiple of res.
        """
        interpolant = interpolant or DEFAULT_INTERPOLANT
        delta = (res[0] / shape[0], res[1] / shape[1])
        d = (shape[0] // res[0], shape[1] // res[1])
        grid = (
            np.mgrid[0 : res[0] : delta[0], 0 : res[1] : delta[1]].transpose(
                1, 2, 0
            )
            % 1
        )
        # Gradients
        angles = 2 * np.pi * np.random.rand(res[0] + 1, res[1] + 1)
        gradients = np.dstack((np.cos(angles), np.sin(angles)))
        if tileable[0]:
            gradients[-1, :] = gradients[0, :]
        if tileable[1]:
            gradients[:, -1] = gradients[:, 0]
        gradients = gradients.repeat(d[0], 0).repeat(d[1], 1)
        g00 = gradients[: -d[0], : -d[1]]
        g10 = gradients[d[0] :, : -d[1]]
        g01 = gradients[: -d[0], d[1] :]
        g11 = gradients[d[0] :, d[1] :]
        # Ramps
        n00 = np.sum(np.dstack((grid[:, :, 0], grid[:, :, 1])) * g00, 2)
        n10 = np.sum(np.dstack((grid[:, :, 0] - 1, grid[:, :, 1])) * g10, 2)
        n01 = np.sum(np.dstack((grid[:, :, 0], grid[:, :, 1] - 1)) * g01, 2)
        n11 = np.sum(
            np.dstack((grid[:, :, 0] - 1, grid[:, :, 1] - 1)) * g11, 2
        )
        # Interpolation
        t = interpolant(grid)
        n0 = n00 * (1 - t[:, :, 0]) + t[:, :, 0] * n10
        n1 = n01 * (1 - t[:, :, 0]) + t[:, :, 0] * n11
        return np.sqrt(2) * ((1 - t[:, :, 1]) * n0 + t[:, :, 1] * n1)

    def generate_fractal_noise_2d(
        self,
        shape,
        res,
        octaves=1,
        persistence=0.5,
        lacunarity=2,
        tileable=(True, True),
        interpolant=None,
    ):
        """Generate a 2D numpy array of fractal noise.

        Args:
            shape: The shape of the generated array (tuple of two ints).
                This must be a multiple of lacunarity**(octaves-1)*res.
            res: The number of periods of noise to generate along each
                axis (tuple of two ints). Note shape must be a multiple of
                (lacunarity**(octaves-1)*res).
            octaves: The number of octaves in the noise. Defaults to 1.
            persistence: The scaling factor between two octaves.
            lacunarity: The frequency factor between two octaves.
            tileable: If the noise should be tileable along each axis
                (tuple of two bools). Defaults to (True,True).
            interpolant: The, interpolation function, defaults to
                t*t*t*(t*(t*6 - 15) + 10).

        Returns
        -------
            A numpy array of fractal noise and of shape shape generated by
            combining several octaves of perlin noise.

        Raises
        ------
            ValueError: If shape is not a multiple of
                (lacunarity**(octaves-1)*res).
        """
        interpolant = interpolant or DEFAULT_INTERPOLANT

        noise = np.zeros(shape)
        frequency = 1
        amplitude = 1
        for _ in range(octaves):
            noise += amplitude * self.generate_perlin_noise_2d(
                shape,
                (frequency * res[0], frequency * res[1]),
                tileable,
                interpolant,
            )
            frequency *= lacunarity
            amplitude *= persistence
        return noise

    def fbm(self, x, y, octaves):
        # noise_2d = self.generate_fractal_noise_2d((256, 256), (8, 8), octaves)
        # Now, extract a single noise value based on x and y, wrapping indices if necessary
        x_idx = int(x) % 256
        y_idx = int(y) % 256
        return self.noise_pattern[x_idx, y_idx]

    def apply_shake(
        self,
        images,
        position_amount_x,
        position_amount_y,
        rotation_amount,
        frequency,
        frequency_divider,
        octaves,
        seed,
    ):
        # Rehash
        np.random.seed(seed)
        self.position_offset = np.random.uniform(-1e3, 1e3, 3)
        self.rotation_offset = np.random.uniform(-1e3, 1e3, 3)
        self.noise_pattern = self.generate_perlin_noise_2d(
            (512, 512), (32, 32), (True, True)
        )

        # Assuming frame count is derived from the first dimension of images tensor
        frame_count = images.shape[0]

        frequency = frequency / frequency_divider

        # Generate shaking parameters for each frame
        x_translations = []
        y_translations = []
        rotations = []

        for frame_num in range(frame_count):
            time = frame_num * frequency
            x_idx = (self.position_offset[0] + frame_num) % 256
            y_idx = (self.position_offset[1] + frame_num) % 256

            np_position = np.array(
                [
                    self.fbm(x_idx, time, octaves),
                    self.fbm(y_idx, time, octaves),
                ]
            )

            # np_position = np.array(
            #     [
            #         self.fbm(self.position_offset[0] + frame_num, time, octaves),
            #         self.fbm(self.position_offset[1] + frame_num, time, octaves),
            #     ]
            # )
            # np_rotation = self.fbm(self.rotation_offset[2] + frame_num, time, octaves)

            rot_idx = (self.rotation_offset[2] + frame_num) % 256
            np_rotation = self.fbm(rot_idx, time, octaves)

            x_translations.append(np_position[0] * position_amount_x)
            y_translations.append(np_position[1] * position_amount_y)
            rotations.append(np_rotation * rotation_amount)

        # Convert lists to tensors
        # x_translations = torch.tensor(x_translations, dtype=torch.float32)
        # y_translations = torch.tensor(y_translations, dtype=torch.float32)
        # rotations = torch.tensor(rotations, dtype=torch.float32)

        # Create an instance of Batch2dTransform
        transform = MTB_Batch2dTransform()

        log.debug(
            f"Applying shaking with parameters: \nposition {position_amount_x}, {position_amount_y}\nrotation {rotation_amount}\nfrequency {frequency}\noctaves {octaves}"
        )

        # Apply shaking transformations to images
        shaken_images = transform.transform_batch(
            images,
            border_handling="edge",  # Assuming edge handling as default
            constant_color="#000000",  # Assuming black as default constant color
            x=x_translations,
            y=y_translations,
            angle=rotations,
        )[0]

        return (shaken_images, x_translations, y_translations, rotations)


__nodes__ = [
    MTB_BatchFloat,
    MTB_Batch2dTransform,
    MTB_BatchShape,
    MTB_BatchMake,
    MTB_BatchFloatAssemble,
    MTB_BatchFloatFill,
    MTB_BatchFloatNormalize,
    MTB_BatchMerge,
    MTB_BatchShake,
    MTB_PlotBatchFloat,
    MTB_BatchTimeWrap,
    MTB_BatchFloatFit,
    MTB_BatchFloatMath,
]