Add Conv2d circular padding options for models and VAEs

ReadMe.md

@@ -114,10 +114,18 @@ This node crops a 360 degree equirectangular image into a 180 degree equirectang
 
 This node allows you to horizontally or vertically crop a 360 degree equirectangular image that is in stereo format or 3D format, into two separate monoscopic images.
 
-### "Merge Monoscopic into Stereo
+### Merge Monoscopic into Stereo
 
 This node allows you to combine two separate monoscopic 360 degree images into a single stereo or 3D image.
 
+### Apply Circular Padding Model
+
+This node allows you to apply circular padding to the x-axis of all Conv2d layers in a model. This can help with reducing the seam where the left and right sides of the image connect.
+
+### Apply Circular Padding VAE
+
+This node allows you to apply circular padding to the x-axis of all the Conv2d layers in a VAE. This can help with reducing the seam where the left and right sides of the image connect.
+
 ### Panorama Viewer
 
 For viewing 360 images inside ComfyUI, see the [github.com/ProGamerGov/ComfyUI_preview360panorama](https://github.com/ProGamerGov/ComfyUI_preview360panorama) custom node.

__init__.py

@@ -1,4 +1,6 @@
 from .nodes import (
+    ApplyCircularConvPaddingModel,
+    ApplyCircularConvPaddingVAE,
     C2EMaskedDiffNode,
     C2ENode,
     Crop360To180Node,
@@ -30,6 +32,8 @@ NODE_CLASS_MAPPINGS = {
     "Crop 360 to 180 Equirectangular": Crop360To180Node,
     "Crop Stereo to Monoscopic": StereoToMonoScopicNode,
     "Merge Monoscopic into Stereo": StereoToMonoScopicNode,
+    "Apply Circular Padding VAE": ApplyCircularConvPaddingVAE,
+    "Apply Circular Padding Model": ApplyCircularConvPaddingModel,
 }
 
 NODE_DISPLAY_NAME_MAPPINGS = {
@@ -47,4 +51,6 @@ NODE_DISPLAY_NAME_MAPPINGS = {
     "Crop 360 to 180 Equirectangular": "Crop 360 to 180 Equirectangular",
     "Crop Stereo to Monoscopic": "Crop Stereo to Monoscopic",
     "Merge Monoscopic into Stereo": "Merge Monoscopic into Stereo",
+    "Apply Circular Padding VAE": "Apply Circular Padding VAE",
+    "Apply Circular Padding Model": "Apply Circular Padding Model",
 }

nodes.py

@@ -1,3 +1,4 @@
+import copy
 from typing import Dict, List, Tuple
 
 import torch
@@ -752,3 +753,127 @@ class MonoScopicToStereoNode:
             )
 
         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"
+    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"
+
+    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,)