155 lines
4.8 KiB
Python
155 lines
4.8 KiB
Python
import copy
|
|
from typing import Optional
|
|
|
|
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
|
|
|
|
|
|
class SeamlessTile:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"model": ("MODEL",),
|
|
"tiling": (["enable", "x_only", "y_only", "disable"],),
|
|
"copy_model": (["Make a copy", "Modify in place"],),
|
|
},
|
|
}
|
|
|
|
CATEGORY = "conditioning"
|
|
|
|
RETURN_TYPES = ("MODEL",)
|
|
FUNCTION = "run"
|
|
|
|
def run(self, model, copy_model, tiling):
|
|
if copy_model == "Modify in place":
|
|
model_copy = model
|
|
else:
|
|
model_copy = copy.deepcopy(model)
|
|
|
|
if tiling == "enable":
|
|
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:
|
|
make_circular_asymm(model_copy.model, False, False)
|
|
return (model_copy,)
|
|
|
|
|
|
# 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)
|
|
|
|
|
|
class CircularVAEDecode:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"samples": ("LATENT",),
|
|
"vae": ("VAE",),
|
|
"tiling": (["enable", "x_only", "y_only", "disable"],)
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE",)
|
|
FUNCTION = "decode"
|
|
|
|
CATEGORY = "latent"
|
|
|
|
def decode(self, samples, vae, tiling):
|
|
vae_copy = copy.deepcopy(vae)
|
|
|
|
if tiling == "enable":
|
|
make_circular_asymm(vae_copy.first_stage_model, True, True)
|
|
elif tiling == "x_only":
|
|
make_circular_asymm(vae_copy.first_stage_model, True, False)
|
|
elif tiling == "y_only":
|
|
make_circular_asymm(vae_copy.first_stage_model, False, True)
|
|
else:
|
|
make_circular_asymm(vae_copy.first_stage_model, False, False)
|
|
|
|
result = (vae_copy.decode(samples["samples"]),)
|
|
return result
|
|
|
|
|
|
class MakeCircularVAE:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"vae": ("VAE",),
|
|
"tiling": (["enable", "x_only", "y_only", "disable"],),
|
|
"copy_vae": (["Make a copy", "Modify in place"],),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("VAE",)
|
|
FUNCTION = "run"
|
|
CATEGORY = "latent"
|
|
|
|
def run(self, vae, tiling, copy_vae):
|
|
if copy_vae == "Modify in place":
|
|
vae_copy = vae
|
|
else:
|
|
vae_copy = copy.deepcopy(vae)
|
|
|
|
if tiling == "enable":
|
|
make_circular_asymm(vae_copy.first_stage_model, True, True)
|
|
elif tiling == "x_only":
|
|
make_circular_asymm(vae_copy.first_stage_model, True, False)
|
|
elif tiling == "y_only":
|
|
make_circular_asymm(vae_copy.first_stage_model, False, True)
|
|
else:
|
|
make_circular_asymm(vae_copy.first_stage_model, False, False)
|
|
|
|
return (vae_copy,)
|
|
|
|
|
|
class OffsetImage:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"pixels": ("IMAGE",),
|
|
"x_percent": (
|
|
"FLOAT",
|
|
{"default": 50.0, "min": 0.0, "max": 100.0, "step": 1},
|
|
),
|
|
"y_percent": (
|
|
"FLOAT",
|
|
{"default": 50.0, "min": 0.0, "max": 100.0, "step": 1},
|
|
),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE",)
|
|
FUNCTION = "run"
|
|
CATEGORY = "image"
|
|
|
|
def run(self, pixels, x_percent, y_percent):
|
|
n, y, x, c = pixels.size()
|
|
y = round(y * y_percent / 100)
|
|
x = round(x * x_percent / 100)
|
|
return (pixels.roll((y, x), (1, 2)),)
|