diff --git a/SeamlessTile.py b/SeamlessTile.py index 3bd76e5..28c1eb2 100644 --- a/SeamlessTile.py +++ b/SeamlessTile.py @@ -2,6 +2,11 @@ import copy import PIL import torch +from torch import Tensor +from torch.nn import Conv2d +from torch.nn import functional as F +from torch.nn.modules.utils import _pair +from typing import Optional class SeamlessTile: @@ -10,8 +15,8 @@ class SeamlessTile: return { "required": { "model": ("MODEL",), - "tiling": (["enable", "disable"],), - "copy_model": (["Modify in place", "Make a copy"],), + "tiling": (["enable", "x_only", "y_only", "disable"],), + "copy_model": (["Make a copy", "Modify in place"],), }, } @@ -25,10 +30,15 @@ class SeamlessTile: model_copy = model else: model_copy = copy.deepcopy(model) + if tiling == "enable": - model_copy.model.apply(make_circular) + make_circular_asymm(model_copy.model, True, True) + elif tiling == "x_only": + make_circular_asymm(model_copy.model, True, False) + elif tiling == "y_only": + make_circular_asymm(model_copy.model, False, True) else: - model_copy.model.apply(unmake_circular) + make_circular_asymm(model_copy.model, False, False) return (model_copy,) @@ -37,11 +47,39 @@ def make_circular(m): m.padding_mode = "circular" +# asymmetric tiling from https://github.com/tjm35/asymmetric-tiling-sd-webui/blob/main/scripts/asymmetric_tiling.py +def make_circular_asymm(model, tileX: bool, tileY: bool): + for layer in [ + layer for layer in model.modules() if isinstance(layer, torch.nn.Conv2d) + ]: + layer.padding_modeX = 'circular' if tileX else 'constant' + layer.padding_modeY = 'circular' if tileY else 'constant' + layer.paddingX = (layer._reversed_padding_repeated_twice[0], layer._reversed_padding_repeated_twice[1], 0, 0) + layer.paddingY = (0, 0, layer._reversed_padding_repeated_twice[2], layer._reversed_padding_repeated_twice[3]) + layer._conv_forward = __replacementConv2DConvForward.__get__(layer, Conv2d) + return model + + +def __replacementConv2DConvForward(self, input: Tensor, weight: Tensor, bias: Optional[Tensor]): + working = F.pad(input, self.paddingX, mode=self.padding_modeX) + working = F.pad(working, self.paddingY, mode=self.padding_modeY) + return F.conv2d(working, weight, bias, self.stride, _pair(0), self.dilation, self.groups) + + def unmake_circular(m): if isinstance(m, torch.nn.Conv2d): m.padding_mode = "zeros" +def unmake_circular_asymm(model): + for layer in [ + layer for layer in model.modules() if isinstance(layer, torch.nn.Conv2d) + ]: + layer.padding_mode = "zeros" + layer._conv_forward = Conv2d._conv_forward.__get__(layer, Conv2d) + return model + + class CircularVAEDecode: @classmethod def INPUT_TYPES(s):