New mask nodes & reordering of all nodes

ProGamerGov · web-flow · commit 240d08127da6 · 2025-08-21T16:23:28.000-06:00
diff --git a/__init__.py b/__init__.py
@@ -11,7 +11,10 @@
     E2ENode,
     E2Face,
     E2PNode,
+    EMask2Face,
     Face2E,
+    FaceMask2E,
+    MaskE2ENode,
     MonoScopicToStereoNode,
     Pad180To360Node,
     PasteImageWithCoordsNode,
@@ -23,49 +26,69 @@
 )
 
 NODE_CLASS_MAPPINGS = {
+    # Equirectangular Single Face
+    "Equirectangular to Face": E2Face,
+    "Equirectangular Mask to Face": EMask2Face,
+    "Face to Equirectangular": Face2E,
+    "Face Mask to Equirectangular": FaceMask2E,
+    # Equirectangular Full
     "Cubemap to Equirectangular": C2ENode,
     "Equirectangular to Cubemap": E2CNode,
     "Equirectangular to Perspective": E2PNode,
     "Equirectangular Rotation": E2ENode,
+    "Mask Equirectangular Rotation": MaskE2ENode,
+    "Pad 180 to 360 Equirectangular": Pad180To360Node,
+    "Crop 360 to 180 Equirectangular": Crop360To180Node,
+    # Cubemap
     "Split Cubemap Faces": SplitFacesNode,
     "Stack Cubemap Faces": StackFacesNode,
+    # Miscellaneous
     "Roll Image Axes": RollImageNode,
+    "Roll Mask Axes": RollMaskNode,
     "Masked Diff C2E": C2EMaskedDiffNode,
     "Crop Image with Coords": CropImageWithCoordsNode,
     "Paste Image with Coords": PasteImageWithCoordsNode,
-    "Pad 180 to 360 Equirectangular": Pad180To360Node,
-    "Crop 360 to 180 Equirectangular": Crop360To180Node,
+    # Stereo
     "Crop Stereo to Monoscopic": StereoToMonoScopicNode,
     "Merge Monoscopic into Stereo": StereoToMonoScopicNode,
+    # Models
     "Apply Circular Padding VAE": ApplyCircularConvPaddingVAE,
     "Apply Circular Padding Model": ApplyCircularConvPaddingModel,
+    # Masks
     "Create Seam Mask": CreateSeamMask,
     "Create Pole Mask": CreatePoleMask,
-    "Equirectangular to Face": E2Face,
-    "Face to Equirectangular": Face2E,
-    "Roll Mask Axes": RollMaskNode,
 }
 
 NODE_DISPLAY_NAME_MAPPINGS = {
+    # Equirectangular Single Face
+    "Equirectangular to Face": "Equirectangular to Face",
+    "Equirectangular Mask to Face": "Equirectangular Mask to Face",
+    "Face to Equirectangular": "Face to Equirectangular",
+    "Face Mask to Equirectangular": "Face Mask to Equirectangular",
+    # Equirectangular Full
     "Cubemap to Equirectangular": "Cubemap to Equirectangular",
     "Equirectangular to Cubemap": "Equirectangular to Cubemap",
     "Equirectangular to Perspective": "Equirectangular to Perspective",
     "Equirectangular Rotation": "Equirectangular Rotation",
+    "Mask Equirectangular Rotation": "Mask Equirectangular Rotation",
+    "Pad 180 to 360 Equirectangular": "Pad 180 to 360 Equirectangular",
+    "Crop 360 to 180 Equirectangular": "Crop 360 to 180 Equirectangular",
+    # Cubemap
     "Split Cubemap Faces": "Split Cubemap Faces",
     "Stack Cubemap Faces": "Stack Cubemap Faces",
+    # Miscellaneous
     "Roll Image Axes": "Roll Image Axes",
+    "Roll Mask Axes": "Roll Mask Axes",
     "Masked Diff C2E": "Masked Diff C2E",
     "Crop Image with Coords": "Crop Image with Coords",
     "Paste Image with Coords": "Paste Image with Coords",
-    "Pad 180 to 360 Equirectangular": "Pad 180 to 360 Equirectangular",
-    "Crop 360 to 180 Equirectangular": "Crop 360 to 180 Equirectangular",
+    # Stereo
     "Crop Stereo to Monoscopic": "Crop Stereo to Monoscopic",
     "Merge Monoscopic into Stereo": "Merge Monoscopic into Stereo",
+    # Models
     "Apply Circular Padding VAE": "Apply Circular Padding VAE",
     "Apply Circular Padding Model": "Apply Circular Padding Model",
+    # Masks
     "Create Seam Mask": "Create Seam Mask",
     "Create Pole Mask": "Create Pole Mask",
-    "Equirectangular to Face": "Equirectangular to Face",
-    "Face to Equirectangular": "Face to Equirectangular",
-    "Roll Mask Axes": "Roll Mask Axes",
 }
diff --git a/nodes.py b/nodes.py
@@ -1194,7 +1194,7 @@ def run(
             # 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:1, ...],)
+        return (output_mask[:, 0, ...],)
 
 
 class Face2E:
@@ -1299,3 +1299,173 @@ def roll_mask(
             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"
+
+    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"
+
+    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"
+
+    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],
+        )
diff --git a/pyproject.toml b/pyproject.toml
@@ -1,7 +1,7 @@
 [project]
 name = "comfyui-pytorch360convert"
 description = "A collection of custom nodes for working with and converting between and working with 360 degree equirectangular images, cubemap, and perspective images. Also includes Circular padding Conv2d options for models and VAEs, along with nodes to fix seams and polar artifacts. Panoramic 360 images are also sometimes known as VR photography (virtual reality), HDRI environments (ex: skyboxes), image spheres, spherical images, 360 pano."
-version = "1.0.9"
+version = "1.1.0"
 license = {file = "LICENSE"}
 dependencies = ["pytorch360convert"]
 
