diff --git a/ReadMe.md b/ReadMe.md index 2f0fe17..9eeab1e 100644 --- a/ReadMe.md +++ b/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. diff --git a/__init__.py b/__init__.py index 85fafc6..cee02d3 100644 --- a/__init__.py +++ b/__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", } diff --git a/nodes.py b/nodes.py index fdb3b76..df62aa9 100644 --- a/nodes.py +++ b/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,)