nodes.py

import copy
from typing import Dict, List, Tuple

import torch
from pytorch360convert import c2e, e2c, e2e, e2p


class C2ENode:
    """
    Cubemap To Equirectangular Node
    """

    @classmethod
    def INPUT_TYPES(s) -> Dict:
        return {
            "required": {
                "e_img": ("IMAGE", {"default": None}),
                "h": ("INT", {"default": -1}),
                "w": ("INT", {"default": -1}),
                "padding_mode": (
                    ["bilinear", "bicubic", "nearest"],
                    {"default": "bilinear"},
                ),
                "cube_format": (
                    ["stack", "dice", "horizon", "list", "dict"],
                    {"default": "stack"},
                ),
            },
        }

    RETURN_TYPES = ("IMAGE",)
    RETURN_NAMES = ("Equirectangular Image",)

    FUNCTION = "c2e"

    CATEGORY = "pytorch360convert/equirectangular"

    def c2e(
        self,
        e_img: torch.Tensor,
        h=None,
        w=None,
        padding_mode: str = "bilinear",
        cube_format: str = "stack",
    ) -> Tuple[torch.Tensor]:
        assert e_img.shape[0] == 6, f"Input should have 6 faces, got {e_img.shape[0]}"
        h = None if h < 1 else h
        w = None if w < 1 else w
        return (
            c2e(
                e_img.squeeze(0),
                h=h,
                w=w,
                cube_format=cube_format,
                mode=padding_mode,
                channels_first=False,
            ).unsqueeze(0),
        )


class E2CNode:
    """
    Equirectangular to Cubemap Node
    """

    @classmethod
    def INPUT_TYPES(s) -> Dict:
        return {
            "required": {
                "e_img": ("IMAGE", {"default": None}),
                "face_width": ("INT", {"default": -1}),
                "padding_mode": (
                    ["bilinear", "bicubic", "nearest"],
                    {"default": "bilinear"},
                ),
                "cube_format": (
                    ["stack", "dice", "horizon", "list", "dict"],
                    {"default": "stack"},
                ),
            },
        }

    RETURN_TYPES = ("IMAGE",)
    RETURN_NAMES = ("Cubemap Image",)

    FUNCTION = "e2c"

    CATEGORY = "pytorch360convert/equirectangular"

    def e2c(
        self,
        e_img: torch.Tensor,
        face_width: int = -1,
        padding_mode: str = "bilinear",
        cube_format: str = "stack",
    ) -> Tuple[torch.Tensor]:
        assert e_img.shape[0] == 1, (
            "Only a batch size of 1 is currently" + f"supported, got {e_img.shape[0]}"
        )
        face_width = e_img.shape[1] // 2 if face_width < 1 else face_width
        output = e2c(
            e_img.squeeze(0),
            face_w=face_width,
            mode=padding_mode,
            cube_format=cube_format,
            channels_first=False,
        )
        output = output.unsqueeze(0) if output.dim() == 3 else output
        return (output,)


class E2PNode:
    """
    Equirectangular to Perspective Node
    """

    @classmethod
    def INPUT_TYPES(s) -> Dict:
        return {
            "required": {
                "e_img": ("IMAGE", {"default": None}),
                "fov_deg_h": ("FLOAT", {"default": 90.0}),
                "fov_deg_v": ("FLOAT", {"default": 90.0}),
                "h_deg": ("FLOAT", {"default": 0.0}),
                "v_deg": ("FLOAT", {"default": 0.0}),
                "out_h": ("INT", {"default": 512}),
                "out_w": ("INT", {"default": 512}),
                "in_rot_deg": ("FLOAT", {"default": 0.0}),
                "padding_mode": (
                    ["bilinear", "bicubic", "nearest"],
                    {"default": "bilinear"},
                ),
            },
        }

    RETURN_TYPES = ("IMAGE",)
    RETURN_NAMES = ("Perspective Image",)

    FUNCTION = "e2p"

    CATEGORY = "pytorch360convert/equirectangular"

    def e2p(
        self,
        e_img: torch.Tensor,
        fov_deg_h: float = 90.0,
        fov_deg_v: float = 90.0,
        h_deg: float = 0.0,
        v_deg: float = 0.0,
        out_h: int = 512,
        out_w: int = 512,
        in_rot_deg: float = 0.0,
        padding_mode: str = "bilinear",
    ) -> Tuple[torch.Tensor]:
        assert e_img.dim() == 4, f"e_img should have 4 dimensions, got {e_img.dim()}"
        return (
            e2p(
                e_img,
                fov_deg=(fov_deg_h, fov_deg_v),
                h_deg=h_deg,
                v_deg=v_deg,
                out_hw=(out_h, out_w),
                in_rot_deg=in_rot_deg,
                mode=padding_mode,
                channels_first=False,
            ),
        )


class E2ENode:
    """
    Equirectangular Rotation Node
    """

    @classmethod
    def INPUT_TYPES(s) -> Dict:
        return {
            "required": {
                "e_img": ("IMAGE", {"default": None}),
                "roll": ("FLOAT", {"default": 0.0}),
                "h_deg": ("FLOAT", {"default": 0.0}),
                "v_deg": ("FLOAT", {"default": 0.0}),
                "padding_mode": (
                    ["bilinear", "bicubic", "nearest"],
                    {"default": "bilinear"},
                ),
            },
        }

    RETURN_TYPES = ("IMAGE",)
    RETURN_NAMES = ("Rotated Image",)

    FUNCTION = "e2e"

    CATEGORY = "pytorch360convert/equirectangular"

    def e2e(
        self,
        e_img: torch.Tensor,
        roll: float = 0.0,
        h_deg: float = 0.0,
        v_deg: float = 0.0,
        padding_mode: str = "bilinear",
    ) -> Tuple[torch.Tensor]:
        assert e_img.dim() == 4, f"e_img should have 4 dimensions, got {e_img.dim()}"
        return (
            e2e(
                e_img=e_img,
                h_deg=h_deg,
                v_deg=v_deg,
                roll=roll,
                mode=padding_mode,
                channels_first=False,
            ),
        )


class SplitFacesNode:
    """
    Split a stack of cube faces for easy manipulation.
    """

    @classmethod
    def INPUT_TYPES(s) -> Dict:
        return {
            "required": {
                "face_stack": ("IMAGE", {"default": None}),
            },
        }

    RETURN_TYPES = ("IMAGE", "IMAGE", "IMAGE", "IMAGE", "IMAGE", "IMAGE")
    RETURN_NAMES = ("Front", "Right", "Back", "Left", "Up", "Down")

    FUNCTION = "split_faces"

    CATEGORY = "pytorch360convert/miscellaneous"

    def split_faces(self, face_stack: torch.Tensor) -> Tuple[torch.Tensor, ...]:
        assert (
            face_stack.dim() == 4
        ), f"face_stack should have 4 dimensions, got {face_stack.dim()}"
        assert (
            face_stack.shape[0] == 6
        ), f"face_stack should have 6 faces, got {face_stack.shape[0]}"
        outputs = [face_stack[i].unsqueeze(0) for i in range(face_stack.shape[0])]
        return tuple(outputs)


class StackFacesNode:
    """
    Combine multiple cube faces into a stack.
    """

    @classmethod
    def INPUT_TYPES(s) -> Dict:
        return {
            "required": {
                "Front": ("IMAGE", {"default": None}),
                "Right": ("IMAGE", {"default": None}),
                "Back": ("IMAGE", {"default": None}),
                "Left": ("IMAGE", {"default": None}),
                "Up": ("IMAGE", {"default": None}),
                "Down": ("IMAGE", {"default": None}),
            },
        }

    RETURN_TYPES = ("IMAGE",)
    RETURN_NAMES = ("Cubemap stack",)

    FUNCTION = "stack_faces"

    CATEGORY = "pytorch360convert/miscellaneous"

    def stack_faces(
        self,
        Front: torch.Tensor,
        Right: torch.Tensor,
        Back: torch.Tensor,
        Left: torch.Tensor,
        Up: torch.Tensor,
        Down: torch.Tensor,
    ) -> Tuple[torch.Tensor]:
        assert (
            Front.dim() == 4
            and Right.dim() == 4
            and Back.dim() == 4
            and Left.dim() == 4
            and Up.dim() == 4
            and Down.dim() == 4
        )
        return (torch.cat([Front, Right, Back, Left, Up, Down], 0),)


class RollImageNode:
    """
    Roll an image to make face seams easier to remove or access.
    """

    @classmethod
    def INPUT_TYPES(s) -> Dict:
        return {
            "required": {
                "image": ("IMAGE", {"default": None}),
                "roll_x": ("INT", {"default": 0}),
                "roll_y": ("INT", {"default": 0}),
                "roll_x_by_50_percent": (
                    "BOOLEAN",
                    {
                        "default": False,
                        "tooltip": "Ignores roll_x and roll_y. Shifts image horizontally by 50%.",
                    },
                ),
            },
        }

    RETURN_TYPES = ("IMAGE",)
    RETURN_NAMES = ("Rolled Image",)

    FUNCTION = "roll_image"

    CATEGORY = "pytorch360convert/miscellaneous"

    def roll_image(
        self,
        image: torch.Tensor,
        roll_x: int = 0,
        roll_y: int = 0,
        roll_x_by_50_percent: bool = False,
    ) -> Tuple[torch.Tensor]:
        assert image.dim() == 4, f"image should have 4 dimensions, got {image.dim()}"
        if roll_x_by_50_percent:
            _, H, W, _ = image.shape
            px_half = W // 2
            roll_y = 0
            roll_x = px_half
        return (torch.roll(image, shifts=(roll_y, roll_x), dims=(1, 2)),)


class C2EMaskedDiffNode:
    """
    Cubemap To Equirectangular with Masked Diff Node
    """

    @classmethod
    def INPUT_TYPES(s) -> Dict:
        return {
            "required": {
                "original_faces": ("IMAGE", {"default": None}),
                "modified_faces": ("IMAGE", {"default": None}),
                "original_equi": ("IMAGE", {"default": None}),
                "padding_mode": (
                    ["bilinear", "bicubic", "nearest"],
                    {"default": "bilinear"},
                ),
                "cube_format": (
                    ["stack", "dice", "horizon", "list", "dict"],
                    {"default": "stack"},
                ),
            },
        }

    RETURN_TYPES = ("IMAGE",)
    RETURN_NAMES = ("Equirectangular Image",)

    FUNCTION = "c2e_masked_diff"

    CATEGORY = "pytorch360convert/mask"

    def c2e_masked_diff(
        self,
        original_faces: torch.Tensor,
        modified_faces: torch.Tensor,
        original_equi: torch.Tensor,
        padding_mode: str = "bilinear",
        cube_format: str = "stack",
    ) -> Tuple[torch.Tensor]:
        assert original_equi.shape[0] == 1, (
            "Only a batch size of 1 is currently supported"
            + f"for original_equi, got {original_equi.shape[0]}"
        )
        assert (
            original_faces.shape[0] == 6
        ), f"original_faces should have 6 faces, got {original_faces.shape[0]}"
        assert (
            modified_faces.shape[0] == 6
        ), f"modified_faces should have 6 faces, got {modified_faces.shape[0]}"
        new_equi = c2e(modified_faces, cube_format=cube_format, channels_first=False)

        faces_mask = original_faces != modified_faces
        mask_equi = c2e(
            faces_mask.float(), cube_format=cube_format, channels_first=False
        ).bool()
        return (torch.where(mask_equi, new_equi, original_equi),)


class CropImageWithCoordsNode:
    """
    Crop an section of an image for manipulation
    """

    @classmethod
    def INPUT_TYPES(s) -> Dict:
        return {
            "required": {
                "image": ("IMAGE", {"default": None}),
                "crop_h": ("INT", {"default": 0}),
                "crop_w": ("INT", {"default": 0}),
                "crop_h2": ("INT", {"default": -1}),
                "crop_w2": ("INT", {"default": -1}),
            },
        }

    RETURN_TYPES = (
        "IMAGE",
        "LIST",
    )
    RETURN_NAMES = (
        "Cropped Image",
        "Coords",
    )

    FUNCTION = "crop_image"

    CATEGORY = "pytorch360convert/miscellaneous"

    def crop_image(
        self,
        image: torch.Tensor,
        crop_h: int = 0,
        crop_w: int = 0,
        crop_h2: int = -1,
        crop_w2: int = -1,
    ) -> Tuple[torch.Tensor, List[int]]:
        assert image.dim() == 4, f"image should have 4 dimensions, got {image.dim()}"

        _, H, W, _ = image.shape

        # Calculate the center crop indices
        if crop_h2 < 1:
            start_h = (H - crop_h) // 2
            end_h = start_h + crop_h
        else:
            start_h = crop_h
            end_h = crop_h2

        if crop_w2 < 1:
            start_w = (W - crop_w) // 2
            end_w = start_w + crop_w
        else:
            start_w = crop_w
            end_w = crop_w2

        cropped_tensor = image[:, start_h:end_h, start_w:end_w, :]

        coords = [start_h, end_h, start_w, end_w]
        return (
            cropped_tensor,
            coords,
        )


class PasteImageWithCoordsNode:
    """
    Add a cropped section back to the original image.
    """

    @classmethod
    def INPUT_TYPES(s) -> Dict:
        return {
            "required": {
                "full_image": ("IMAGE", {"default": None}),
                "cropped_image": ("IMAGE", {"default": None}),
                "coords": ("LIST", {"default": None}),
            },
        }

    RETURN_TYPES = ("IMAGE",)
    RETURN_NAMES = ("Pasted Image",)

    FUNCTION = "paste_image"

    CATEGORY = "pytorch360convert/miscellaneous"

    def paste_image(
        self, full_image: torch.Tensor, cropped_image: torch.Tensor, coords: List[int]
    ) -> Tuple[torch.Tensor]:
        assert (
            full_image.dim() == 4
        ), f"full_image should have 4 dimensions, got {full_image.dim()}"
        assert (
            cropped_image.dim() == 4
        ), f"cropped_image should have 4 dimensions, got {cropped_image.dim()}"
        assert len(coords) == 4
        start_h, end_h, start_w, end_w = coords
        full_image[:, start_h:end_h, start_w:end_w, :] = cropped_image
        return (full_image,)


class Pad180To360Node:
    """
    Apply padding to a 180 degree equirectangular image, to make it 360 degrees.
    """

    @classmethod
    def INPUT_TYPES(s) -> Dict:
        return {
            "required": {
                "image": ("IMAGE", {"default": None}),
                "fill_value": ("FLOAT", {"default": 0.0}),
            },
        }

    RETURN_TYPES = ("IMAGE",)
    RETURN_NAMES = ("Padded 360 Image",)

    FUNCTION = "pad_180_to_360_image"

    CATEGORY = "pytorch360convert/equirectangular"

    def pad_180_to_360_image(
        self, image: torch.Tensor, fill_value: float = 0.0
    ) -> Tuple[torch.Tensor]:
        assert image.dim() == 4, f"image should have 4 dimensions, got {image.dim()}"
        image = image.permute(0, 3, 1, 2)
        H, W = image.shape[2:]
        pad_left = W // 2
        pad_right = W - pad_left

        image_padded = torch.nn.functional.pad(
            image, (pad_left, pad_right), mode="constant", value=fill_value
        )

        image_padded = image_padded.permute(0, 2, 3, 1)
        return (image_padded,)


class Crop360To180Node:
    """
    Crop a 360 degree equirectangular image to the central 180 degree part.
    """

    @classmethod
    def INPUT_TYPES(s) -> Dict:
        return {
            "required": {
                "image": ("IMAGE", {"default": None}),
            },
        }

    RETURN_TYPES = ("IMAGE",)
    RETURN_NAMES = ("Cropped 180 Image",)

    FUNCTION = "crop_360_to_180_image"

    CATEGORY = "pytorch360convert/equirectangular"

    def crop_360_to_180_image(self, image: torch.Tensor) -> Tuple[torch.Tensor]:
        """
        Crop a 360-degree equirectangular image to the central 180-degree part.

        Args:
            image (torch.Tensor): The 360-degree equirectangular image. Shape should
                be: (B, H, W, C).

        Returns:
            torch.Tensor: The cropped 180-degree equirectangular image. Shape will be:
                (B, H, W//2, C) where the width is halved.

        Raises:
            ValueError: If the input image width is less than or equal to 1.
        """
        assert image.dim() == 4, f"image should have 4 dimensions, got {image.dim()}"
        _, _, width, _ = image.shape

        # Crop the central 180-degree part by slicing the image.
        crop_start = width // 4
        crop_end = 3 * width // 4

        # Crop along the width dimension (left and right 180 degrees)
        cropped_img = image[:, :, crop_start:crop_end, :]

        return (cropped_img,)


class StereoToMonoScopicNode:
    """
    Split a stereo image into 2 monoscopic images.
    """

    @classmethod
    def INPUT_TYPES(s) -> Dict:
        return {
            "required": {
                "image": ("IMAGE", {"default": None}),
                "split_direction": (
                    ["horizontal", "vertical"],
                    {"default": "horizontal"},
                ),
                "larger_side": (["first", "second"], {"default": "first"}),
            },
        }

    RETURN_TYPES = (
        "IMAGE",
        "IMAGE",
    )
    RETURN_NAMES = (
        "First Image",
        "Second Image",
    )

    FUNCTION = "split_stereo_image"

    CATEGORY = "pytorch360convert/stereo"

    def split_stereo_image(
        self,
        image: torch.Tensor,
        split_direction: str = "horizontal",
        larger_side: str = "first",
    ) -> Tuple[torch.Tensor, torch.Tensor]:
        """
        Splits a batch of tensor images (BHWC format) into two halves either vertically
        or horizontally,  splitting exactly down the middle. In case of odd dimensions,
        one side gets an extra pixel.

        Args:
            image (torch.Tensor): The tensor image batch to split. Shape should be:
                (B, H, W, C).
            split_direction (str, optional): The direction to split the image. Either
                "horizontal" or "vertical". Default: "horizontal"
            larger_side (str, optional): Specifies which side gets the extra pixel when
                dimensions are odd. "first" gives the extra pixel to the first side
                (top/left), "second" gives it to the second side (bottom/right).
                Default: "first"

        Returns:
            Tuple[torch.Tensor, torch.Tensor]: A tuple containing the first and second
                halves of the split images. The shape of each tensor is:
                - For horizontal split: (B, H/2, W, C)
                - For vertical split: (B, H, W/2, C)

        Raises:
            ValueError: If `split_direction` is not "horizontal" or "vertical".
        """
        _, height, width, channels = image.shape

        if split_direction == "vertical":
            # Split horizontally (by height)
            mid_point = height // 2
            if height % 2 != 0:  # If height is odd
                if larger_side == "first":
                    first_half = image[:, : mid_point + 1, :, :]
                    second_half = image[:, mid_point + 1 :, :, :]
                else:  # "second"
                    first_half = image[:, :mid_point, :, :]
                    second_half = image[:, mid_point:, :]
            else:
                # Even height, just split in the middle
                first_half = image[:, :mid_point, :, :]
                second_half = image[:, mid_point:, :]
        elif split_direction == "horizontal":
            # Split vertically (by width)
            mid_point = width // 2
            if width % 2 != 0:  # If width is odd
                if larger_side == "first":
                    first_half = image[:, :, : mid_point + 1, :]
                    second_half = image[:, :, mid_point + 1 :, :]
                else:  # "second"
                    first_half = image[:, :, :mid_point, :]
                    second_half = image[:, :, mid_point:, :]
            else:
                # Even width, just split in the middle
                first_half = image[:, :, :mid_point, :]
                second_half = image[:, :, mid_point:, :]
        else:
            raise ValueError(
                "Invalid split direction. Please choose"
                + " 'horizontal' or 'vertical'. "
                + f"Got {split_direction}"
            )

        return (first_half, second_half)


class MonoScopicToStereoNode:
    """
    Merge two monoscopic images into a stereo image.
    """

    @classmethod
    def INPUT_TYPES(s) -> Dict:
        return {
            "required": {
                "first_image": ("IMAGE", {"default": None}),
                "second_image": ("IMAGE", {"default": None}),
                "merge_direction": (
                    ["horizontal", "vertical"],
                    {"default": "horizontal"},
                ),
            },
        }

    RETURN_TYPES = ("IMAGE",)
    RETURN_NAMES = ("Stereo Image",)

    FUNCTION = "merge_monoscopic_to_stereo"

    CATEGORY = "pytorch360convert/stereo"

    def merge_monoscopic_to_stereo(
        self,
        first_image: torch.Tensor,
        second_image: torch.Tensor,
        merge_direction: str = "horizontal",
    ) -> Tuple[torch.Tensor]:
        """
        Merges two monoscopic images into a single stereo image by concatenating
        the two images either horizontally or vertically.

        Args:
            first_image (torch.Tensor): The left monoscopic image. Shape should
                be: (B, H, W, C).
            second_image (torch.Tensor): The right monoscopic image. Shape should
                be: (B, H, W, C).
            merge_direction (str, optional): The direction to merge the images. Either
                "horizontal" or "vertical". Default: "horizontal"

        Returns:
            torch.Tensor: The merged stereo image.

        Raises:
            ValueError: If `merge_direction` is not "horizontal" or "vertical".
        """

        f_batch_size, f_height, f_width, f_channels = first_image.shape
        s_batch_size, s_height, s_width, s_channels = second_image.shape

        if merge_direction == "horizontal":
            assert f_batch_size == s_batch_size, (
                "Batch size must match: " + f"{f_batch_size} vs {s_batch_size}"
            )
            assert f_height == s_height, (
                "Height must match: " + f"{f_height} vs {s_height}"
            )
            assert f_channels == s_channels, (
                "Channels must match: " + f"{f_channels} vs {s_channels}"
            )

            # Concatenate horizontally (by width)
            stereo_image = torch.cat((first_image, second_image), dim=2)

        elif merge_direction == "vertical":
            assert f_batch_size == s_batch_size, (
                "Batch size must match: " + f"{f_batch_size} vs {s_batch_size}"
            )
            assert f_width == s_width, "Width must match: " + f"{f_width} vs {s_width}"
            assert f_channels == s_channels, (
                "Channels must match: " + f"{f_channels} vs {s_channels}"
            )

            # Concatenate vertically (by height)
            stereo_image = torch.cat((first_image, second_image), dim=1)

        else:
            raise ValueError(
                "Invalid split direction. Please choose"
                + " 'horizontal' or 'vertical'. "
                + f"Got {merge_direction}"
            )

        return (stereo_image,)


def _conv_forward(
    self, x: torch.Tensor, weight: torch.Tensor, bias: torch.Tensor
) -> torch.Tensor:
    x = torch.nn.functional.pad(x, self.padding_values_x, mode="circular")
    x = torch.nn.functional.pad(x, self.padding_values_y, mode="constant")
    return torch.nn.functional.conv2d(
        x, weight, bias, self.stride, (0, 0), self.dilation, self.groups
    )


def _apply_circular_conv2d_padding(
    model: torch.nn.Module, is_vae: bool = False, x_axis_only: bool = True
) -> torch.nn.Module:
    for layer in model.first_stage_model.modules() if is_vae else model.modules():
        if isinstance(layer, torch.nn.Conv2d):
            if x_axis_only:
                layer.padding_values_x = (
                    layer._reversed_padding_repeated_twice[0],
                    layer._reversed_padding_repeated_twice[1],
                    0,
                    0,
                )
                layer.padding_values_y = (
                    0,
                    0,
                    layer._reversed_padding_repeated_twice[2],
                    layer._reversed_padding_repeated_twice[3],
                )
                layer._conv_forward = _conv_forward.__get__(layer, torch.nn.Conv2d)
            else:
                layer.padding_mode = "circular"
    return model


class ApplyCircularConvPaddingModel:
    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                "model": (
                    "MODEL",
                    {"tooltip": "Model to add circular x-axis conv2d padding to."},
                ),
                "inplace": (
                    "BOOLEAN",
                    {
                        "default": True,
                        "tooltip": "Modify "
                        + "the already loaded model (True) or a copy of the model (False). "
                        + "If True, model will have to be reloaded to restore padding to "
                        + "the original values. Modifying inplace will use less memory.",
                    },
                ),
                "x_axis_only": (
                    "BOOLEAN",
                    {
                        "default": True,
                        "tooltip": "Apply"
                        + " circular padding only to the x-axis or to both the x and y axes.",
                    },
                ),
            },
        }

    CATEGORY = "pytorch360convert/models"
    RETURN_TYPES = ("MODEL",)
    FUNCTION = "run"

    def run(
        self, model: torch.nn.Module, inplace: bool = True, x_axis_only: bool = True
    ) -> Tuple[torch.nn.Module]:
        if inplace:
            use_model = model
        else:
            use_model = copy.deepcopy(model)
        _apply_circular_conv2d_padding(use_model.model, False, x_axis_only)
        return (use_model,)


class ApplyCircularConvPaddingVAE:
    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "vae": (
                    "VAE",
                    {"tooltip": "VAE to add circular x-axis conv2d padding to."},
                ),
                "inplace": (
                    "BOOLEAN",
                    {
                        "default": True,
                        "tooltip": "Modify "
                        + "the already loaded VAE (True) or a copy of the VAE (False). "
                        + "If True, VAE will have to be reloaded to restore padding to "
                        + "the original values. Modifying inplace will use less memory.",
                    },
                ),
                "x_axis_only": (
                    "BOOLEAN",
                    {
                        "default": True,
                        "tooltip": "Apply"
                        + " circular padding only to the x-axis or to both the x and y axes.",
                    },
                ),
            }
        }

    RETURN_TYPES = ("VAE",)
    FUNCTION = "run"
    CATEGORY = "pytorch360convert/models"

    def run(
        self, vae: torch.nn.Module, inplace: bool = True, x_axis_only: bool = True
    ) -> Tuple[torch.nn.Module]:
        if inplace:
            use_vae = vae
        else:
            use_vae = copy.deepcopy(vae)
        _apply_circular_conv2d_padding(use_vae, True, x_axis_only)
        return (use_vae,)


def _create_center_seam_mask(
    x: torch.Tensor, frac_width: float = 0.10, pixel_width: int = 0, feather: int = 0
) -> torch.Tensor:
    """
    For a ComfyUI-style mask: shape [B, H, W], values 0 or 1, with optional feathering.

    Args:
        x (torch.Tensor): input tensor with shape [B, H, W, C].
        frac_width (float, optional): fraction of input width for the vertical strip.
        pixel_width (int, optional): width of the seam in pixels (if > 0).
        feather (int, optional): pixel size of feathering on both sides of the mask.

    Returns:
        mask: torch.Tensor of shape [B, H, W] with float values 0.0 to 1.0.
    """
    # Extract batch, height, and width from x
    B, H, W, *_ = x.shape

    # Determine strip width
    if pixel_width > 0:
        strip = pixel_width
    else:
        strip = max(1, int(W * frac_width))

    x0 = (W - strip) // 2
    x1 = x0 + strip

    # Create the mask with zeros
    mask = torch.zeros((B, H, W), dtype=x.dtype, device=x.device)

    if feather <= 0:
        mask[:, :, x0:x1] = 1.0
    else:
        # Create feathered mask
        # Left feather region
        left_feather_start = max(0, x0 - feather)
        left_feather_end = x0
        if left_feather_end > left_feather_start:
            feather_steps = torch.linspace(
                0.0,
                1.0,
                left_feather_end - left_feather_start,
                dtype=x.dtype,
                device=x.device,
            )
            mask[:, :, left_feather_start:left_feather_end] = feather_steps[
                None, None, :
            ]

        # Center region (full mask)
        mask[:, :, x0:x1] = 1.0

        # Right feather region
        right_feather_start = x1
        right_feather_end = min(W, x1 + feather)
        if right_feather_end > right_feather_start:
            feather_steps = torch.linspace(
                1.0,
                0.0,
                right_feather_end - right_feather_start,
                dtype=x.dtype,
                device=x.device,
            )
            mask[:, :, right_feather_start:right_feather_end] = feather_steps[
                None, None, :
            ]

    return mask


class CreateSeamMask:
    """
    Create a seam mask for inpainting on equirectangular images.
    """

    @classmethod
    def INPUT_TYPES(s) -> Dict:
        return {
            "required": {
                "image": ("IMAGE", {"default": None}),
                "frac_width": ("FLOAT", {"default": 0.10}),
                "pixel_width": (
                    "INT",
                    {
                        "default": 0,
                        "tooltip": "Width of the seam in pixels."
                        + " If > 0, frac_width is ignored in favour of width_pixel.",
                    },
                ),
                "feather": (
                    "INT",
                    {"default": 0, "tooltip": "Pixel size of the feathering."},
                ),
                "roll_x_by_50_percent": (
                    "BOOLEAN",
                    {
                        "default": False,
                        "tooltip": "Shifts the output mask horizontally by 50%."
                        + " Equivalent to a 180 degree rotation on an equirectangular image.",
                    },
                ),
            },
        }

    RETURN_TYPES = ("MASK",)
    RETURN_NAMES = ("Seam Mask",)

    FUNCTION = "run"

    CATEGORY = "pytorch360convert/mask"

    def run(
        self,
        image: torch.Tensor,
        frac_width: float = 0.10,
        pixel_width: int = 0,
        feather: int = 0,
        roll_x_by_50_percent: bool = False,
    ) -> Tuple[torch.Tensor]:
        assert image.dim() == 4, "Image should have 4 dimensions"
        _, H, W, _ = image.shape
        px_half = W // 2
        seam_mask = _create_center_seam_mask(
            image, frac_width=frac_width, pixel_width=pixel_width, feather=feather
        )

        if roll_x_by_50_percent:
            seam_mask = torch.roll(seam_mask, shifts=(0, px_half), dims=(1, 2))
        return (seam_mask,)


class E2Face:
    """
    Equirectangular to Face
    """

    @classmethod
    def INPUT_TYPES(s) -> Dict:
        return {
            "required": {
                "e_img": ("IMAGE", {"default": None}),
                "face_width": ("INT", {"default": -1}),
                "padding_mode": (
                    ["bilinear", "bicubic", "nearest"],
                    {"default": "bilinear"},
                ),
                "cube_face": (
                    ["Up", "Down", "Right", "Left", "Front", "Back"],
                    {"default": "Front"},
                ),
            },
        }

    RETURN_TYPES = ("IMAGE",)
    RETURN_NAMES = ("Face Image",)

    FUNCTION = "run_e2face"

    CATEGORY = "pytorch360convert/equirectangular"

    def run_e2face(
        self,
        e_img: torch.Tensor,
        face_width: int = -1,
        padding_mode: str = "bilinear",
        cube_face: str = "Front",
    ) -> Tuple[torch.Tensor]:

        B, H, W, C = e_img.shape
        outputs = []

        for i in range(B):
            single_img = e_img[i]  # [H,W,C]
            # Determine face width
            face_w = H // 2 if face_width < 1 else face_width

            # Convert single equirectangular image to cubemap dict
            cubemap = e2c(
                single_img,
                face_w=face_w,
                mode=padding_mode,
                cube_format="dict",
                channels_first=False,
            )

            # Pick requested face
            face_tensor = cubemap[cube_face]  # [face_w, face_w, C]

            outputs.append(face_tensor.unsqueeze(0))  # [1,H,W,C]

        # Concatenate into [B,H,W,C]
        output_batch = torch.cat(outputs, dim=0)
        return (output_batch,)


def _create_centered_circle_mask(
    x: torch.Tensor,
    circle_radius: float,
    feather: int = 0,
    pixel_radius: int = 0,
) -> torch.Tensor:
    """
    Create a centered circle mask with optional feathering.

    Args:
        x (torch.Tensor): Reference tensor with shape (B, C, H, W).
            The mask will copy device, dtype, and batch size from x.
        circle_radius (float): Fraction (0.0–1.0) of max possible radius (min(H,W)/2).
            1.0 = circle touches edges.
        feather (int, optional): Feather width in pixels, extending outward from
            circle_radius. 0 = hard edge, >0 = smooth falloff. Default: 0
        pixel_radius (int, optional): Explicit radius in pixels. If >0, overrides
            circle_radius. Default: 0

    Returns:
        torch.Tensor: Mask tensor of shape (1, C, H, W), values in [0.0, 1.0].
    """
    _, C, H, W = x.shape
    size = min(H, W)

    # Circle radius in pixels
    max_radius = size / 2.0
    if pixel_radius > 0:
        inner_radius = float(pixel_radius)
    else:
        inner_radius = circle_radius * max_radius
    outer_radius = inner_radius + feather

    # Coordinate grid
    yy, xx = torch.meshgrid(
        torch.arange(size, device=x.device),
        torch.arange(size, device=x.device),
        indexing="ij",
    )
    center = size // 2
    dist = torch.sqrt((xx - center) ** 2 + (yy - center) ** 2)

    # Base mask
    mask = torch.zeros_like(dist, dtype=x.dtype, device=x.device)

    # Inside circle = 1.0
    mask[dist <= inner_radius] = 1.0

    # Feather transition
    if feather > 0:
        transition_zone = (dist > inner_radius) & (dist <= outer_radius)
        mask[transition_zone] = 1.0 - (dist[transition_zone] - inner_radius) / feather

    # Expand to BCHW and center in full HxW
    mask_full = torch.zeros((1, C, H, W), dtype=x.dtype, device=x.device)
    y0 = (H - size) // 2
    x0 = (W - size) // 2
    mask_full[:, :, y0 : y0 + size, x0 : x0 + size] = mask

    return mask_full


class CreatePoleMask:
    """
    Create a pole mask for inpainting on equirectangular images.
    """

    @classmethod
    def INPUT_TYPES(s) -> Dict:
        return {
            "required": {
                "image": ("IMAGE", {"default": None}),
                "circle_radius": ("FLOAT", {"default": 0.10}),
                "pixel_radius": (
                    "INT",
                    {
                        "default": 0,
                        "tooltip": "Size of the radius in pixels."
                        + " If > 0, circle_radius is ignored in favour of pixel_radius.",
                    },
                ),
                "feather": ("INT", {"default": 0}),
                "mode": (["face", "equirectangular"], {"default": "face"}),
            },
        }

    RETURN_TYPES = ("MASK",)
    RETURN_NAMES = ("Pole Mask",)

    FUNCTION = "run"

    CATEGORY = "pytorch360convert/mask"

    def run(
        self,
        image: torch.Tensor,
        circle_radius: float = 0.10,
        pixel_radius: int = 0,
        feather: int = 0,
        mode: str = "face",
    ) -> Tuple[torch.Tensor]:
        assert image.dim() == 4, f"Image should have 4 dimensions, got {image.shape}"

        if mode == "face":
            image = image.permute(0, 3, 1, 2)  # BHWC -> BCHW
            output_mask = _create_centered_circle_mask(
                image,
                circle_radius=circle_radius,
                pixel_radius=pixel_radius,
                feather=feather,
            )
            output_mask = output_mask
        else:
            output_mask = []
            for im in image:  # im: [H,W,C]
                # Convert single equirectangular image to cubemap dict
                cubemap = e2c(
                    torch.zeros_like(im), cube_format="dict", channels_first=False
                )

                # Apply circle mask on Up and Down faces
                for face in ["Up", "Down"]:
                    face_tensor = cubemap[face]  # [H_face, W_face, C]
                    face_tensor = face_tensor.unsqueeze(0).permute(0, 3, 1, 2)
                    face_tensor = _create_centered_circle_mask(
                        face_tensor,
                        circle_radius=circle_radius,
                        pixel_radius=pixel_radius,
                        feather=feather,
                    )
                    cubemap[face] = face_tensor.permute(0, 2, 3, 1).squeeze(0)

                # Convert back to equirectangular
                new_equi = c2e(
                    cubemap, cube_format="dict", channels_first=False
                ).unsqueeze(
                    0
                )  # add batch dim
                output_mask.append(new_equi)

            # Stack all results into a single tensor [B,H,W,C]
            output_mask = torch.cat(output_mask).permute(0, 3, 1, 2)

        return (output_mask[:, 0, ...],)


class Face2E:
    """
    Face To Equirectangular Node
    """

    @classmethod
    def INPUT_TYPES(s) -> Dict:
        return {
            "required": {
                "image": ("IMAGE", {"default": None}),
                "face": (
                    ["Up", "Down", "Front", "Right", "Left", "Back"],
                    {"default": "Down"},
                ),
                "base_equi_color": ("FLOAT", {"default": 0.0}),
                "padding_mode": (
                    ["bilinear", "bicubic", "nearest"],
                    {"default": "bilinear"},
                ),
            },
        }

    RETURN_TYPES = ("IMAGE",)
    RETURN_NAMES = ("Equirectangular Image",)

    FUNCTION = "run_face2e"

    CATEGORY = "pytorch360convert/equirectangular"

    def run_face2e(
        self,
        image: torch.Tensor,
        face: str = "Down",
        base_equi_color: float = 0.0,
        padding_mode: str = "bilinear",
    ) -> Tuple[torch.Tensor]:
        assert image.dim() == 4, f"Image should have 4 dimensions, got {image.shape}"

        output_image = []
        for f_image in image:

            cubemap_dict = {}
            for face_name in ["Front", "Right", "Back", "Left", "Up", "Down"]:
                if face_name != face:
                    cubemap_dict[face_name] = torch.ones_like(f_image) * base_equi_color
                else:
                    cubemap_dict[face_name] = f_image
            output_image += [
                c2e(
                    cubemap=cubemap_dict,
                    cube_format="dict",
                    mode=padding_mode,
                    channels_first=False,
                )
            ]
        return (torch.stack(output_image),)


class RollMaskNode:
    """
    Roll a mask to make face seams easier to remove or access.
    """

    @classmethod
    def INPUT_TYPES(s) -> Dict:
        return {
            "required": {
                "mask": ("MASK", {"default": None}),
                "roll_x": ("INT", {"default": 0}),
                "roll_y": ("INT", {"default": 0}),
                "roll_x_by_50_percent": (
                    "BOOLEAN",
                    {
                        "default": False,
                        "tooltip": "Shifts the input mask horizontally by 50%."
                        + " Equivalent to a 180 degree rotation on an equirectangular image.",
                    },
                ),
            },
        }

    RETURN_TYPES = ("MASK",)
    RETURN_NAMES = ("Rolled Mask",)

    FUNCTION = "roll_mask"

    CATEGORY = "pytorch360convert/mask"

    def roll_mask(
        self,
        mask: torch.Tensor,
        roll_x: int = 0,
        roll_y: int = 0,
        roll_x_by_50_percent: bool = False,
    ) -> Tuple[torch.Tensor]:
        assert mask.dim() == 3, f"mask should have 3 dimensions, got {mask.dim()}"
        if roll_x_by_50_percent:
            _, H, W = mask.shape
            px_half = W // 2
            roll_y = 0
            roll_x = px_half
        return (torch.roll(mask, shifts=(roll_y, roll_x), dims=(1, 2)),)


class FaceMask2E:
    """
    Face Mask To Equirectangular
    """

    @classmethod
    def INPUT_TYPES(s) -> Dict:
        return {
            "required": {
                "mask": ("MASK", {"default": None}),
                "face": (
                    ["Up", "Down", "Front", "Right", "Left", "Back"],
                    {"default": "Down"},
                ),
                "base_equi_color": ("FLOAT", {"default": 0.0}),
                "padding_mode": (
                    ["bilinear", "bicubic", "nearest"],
                    {"default": "bilinear"},
                ),
            },
        }

    RETURN_TYPES = ("MASK",)
    RETURN_NAMES = ("Equirectangular Mask",)

    FUNCTION = "run_facemask2e"

    CATEGORY = "pytorch360convert/mask"

    def run_facemask2e(
        self,
        mask: torch.Tensor,
        face: str = "Down",
        base_equi_color: float = 0.0,
        padding_mode: str = "bilinear",
    ) -> Tuple[torch.Tensor]:
        assert mask.dim() == 3, f"Image should have 3 dimensions, got {mask.shape}"

        output_mask = []
        for f_mask in mask:
            f_mask = f_mask = f_mask[None, ...]
            cubemap_dict = {}
            for face_name in ["Front", "Right", "Back", "Left", "Up", "Down"]:
                if face_name != face:
                    cubemap_dict[face_name] = torch.ones_like(f_mask) * base_equi_color
                else:
                    cubemap_dict[face_name] = mask.reshape(1, *mask.shape[1:])
            output_mask += [
                c2e(
                    cubemap=cubemap_dict,
                    cube_format="dict",
                    mode=padding_mode,
                    channels_first=True,
                )
            ]
        return (torch.cat(output_mask),)


class EMask2Face:
    """
    Equirectangular Mask to Face
    """

    @classmethod
    def INPUT_TYPES(s) -> Dict:
        return {
            "required": {
                "mask": ("MASK", {"default": None}),
                "face_width": ("INT", {"default": -1}),
                "padding_mode": (
                    ["bilinear", "bicubic", "nearest"],
                    {"default": "bilinear"},
                ),
                "cube_face": (
                    ["Up", "Down", "Right", "Left", "Front", "Back"],
                    {"default": "Front"},
                ),
            },
        }

    RETURN_TYPES = ("MASK",)
    RETURN_NAMES = ("Face MASK",)

    FUNCTION = "run_emask2face"

    CATEGORY = "pytorch360convert/mask"

    def run_emask2face(
        self,
        mask: torch.Tensor,
        face_width: int = -1,
        padding_mode: str = "bilinear",
        cube_face: str = "Front",
    ) -> Tuple[torch.Tensor]:

        B, H, W = mask.shape
        outputs = []

        for i in range(B):
            singlmask = mask[i : i + 1][..., None]  # [H,W,C]
            # Determine face width
            face_w = H // 2 if face_width < 1 else face_width

            # Convert single equirectangular mask to cubemap dict
            cubemap = e2c(
                singlmask,
                face_w=face_w,
                mode=padding_mode,
                cube_format="dict",
                channels_first=False,
            )

            # Pick requested face
            face_tensor = cubemap[cube_face]  # [face_w, face_w, C]

            outputs.append(face_tensor)  # [1,H,W,C]

        # Concatenate into [B,H,W,C]
        output_batch = torch.cat(outputs, dim=0)[..., 0]
        return (output_batch,)


class MaskE2ENode:
    """
    Mask Equirectangular Rotation
    """

    @classmethod
    def INPUT_TYPES(s) -> Dict:
        return {
            "required": {
                "mask": ("MASK", {"default": None}),
                "roll": ("FLOAT", {"default": 0.0}),
                "h_deg": ("FLOAT", {"default": 0.0}),
                "v_deg": ("FLOAT", {"default": 0.0}),
                "padding_mode": (
                    ["bilinear", "bicubic", "nearest"],
                    {"default": "bilinear"},
                ),
            },
        }

    RETURN_TYPES = ("MASK",)
    RETURN_NAMES = ("Rotated Mask",)

    FUNCTION = "mask_e2e"

    CATEGORY = "pytorch360convert/mask"

    def mask_e2e(
        self,
        mask: torch.Tensor,
        roll: float = 0.0,
        h_deg: float = 0.0,
        v_deg: float = 0.0,
        padding_mode: str = "bilinear",
    ) -> Tuple[torch.Tensor]:
        assert mask.dim() == 3, f"mask should have 3 dimensions, got {mask.dim()}"
        return (
            e2e(
                e_img=mask[..., None],
                h_deg=h_deg,
                v_deg=v_deg,
                roll=roll,
                mode=padding_mode,
                channels_first=False,
            )[..., 0],
        )


class Create180To360Mask:
    @classmethod
    def INPUT_TYPES(s) -> Dict:
        return {
            "required": {
                "image": ("IMAGE", {"default": None}),
                "input_mode": (
                    ["360", "180"],
                    {"default": "180"},
                ),
                "feather": ("INT", {"default": 0}),
            },
        }

    RETURN_TYPES = ("MASK",)
    RETURN_NAMES = ("Equirectangular Mask",)

    FUNCTION = "mask_180_to_360"

    CATEGORY = "pytorch360convert/mask"

    def mask_180_to_360(
        self, image: torch.Tensor, input_mode: str = "180", feather: int = 0
    ) -> Tuple[torch.Tensor]:
        assert image.dim() == 4, f"image should have 4 dimensions, got {image.dim()}"
        _, H, W, _ = image.shape

        # For 360 input, the valid region is only half width
        if input_mode == "360":
            W = W // 2

        pad_left = W // 2
        pad_right = W - pad_left
        total_width = W + pad_left + pad_right  # == 2*W

        # Start with zeros everywhere
        mask = torch.zeros(H, total_width, dtype=image.dtype, device=image.device)

        # Fill the main region with 1.0
        mask[:, pad_left : pad_left + W] = 1.0

        if feather > 0:
            ramp = torch.linspace(
                0, 1, steps=feather + 1, device=image.device, dtype=image.dtype
            )[1:]

            # Left feather (in the padded zero region, approaching the mask)
            mask[:, pad_left - feather : pad_left] = ramp

            # Right feather (in the padded zero region, approaching the mask)
            mask[:, pad_left + W : pad_left + W + feather] = ramp.flip(0)

        # [1, H, W] mask tensor
        return (mask.unsqueeze(0),)