mirror of
https://github.com/invoke-ai/InvokeAI
synced 2024-08-30 20:32:17 +00:00
Quick Seamless Fixes
This commit is contained in:
parent
daa780940b
commit
c3b2a8cb27
@ -28,68 +28,49 @@ def _conv_forward_asymmetric(self, input, weight, bias):
|
|||||||
def set_seamless(model: Union[UNet2DConditionModel, AutoencoderKL], seamless_axes: List[str]):
|
def set_seamless(model: Union[UNet2DConditionModel, AutoencoderKL], seamless_axes: List[str]):
|
||||||
try:
|
try:
|
||||||
to_restore = []
|
to_restore = []
|
||||||
|
skipped_layers = 0
|
||||||
|
skip_second_resnet = True
|
||||||
|
skip_conv2 = True
|
||||||
|
|
||||||
for m_name, m in model.named_modules():
|
for m_name, m in model.named_modules():
|
||||||
if isinstance(model, UNet2DConditionModel):
|
if not isinstance(m, (nn.Conv2d, nn.ConvTranspose2d)):
|
||||||
if ".attentions." in m_name:
|
continue
|
||||||
|
|
||||||
|
if isinstance(model, UNet2DConditionModel) and m_name.startswith("down_blocks.") and ".resnets." in m_name:
|
||||||
|
# down_blocks.1.resnets.1.conv1
|
||||||
|
_, block_num, _, resnet_num, submodule_name = m_name.split(".")
|
||||||
|
block_num = int(block_num)
|
||||||
|
resnet_num = int(resnet_num)
|
||||||
|
|
||||||
|
# if block_num >= seamless_down_blocks:
|
||||||
|
if block_num >= len(model.down_blocks) - skipped_layers:
|
||||||
continue
|
continue
|
||||||
|
|
||||||
if ".resnets." in m_name:
|
if resnet_num > 0 and skip_second_resnet:
|
||||||
if ".conv2" in m_name:
|
|
||||||
continue
|
|
||||||
if ".conv_shortcut" in m_name:
|
|
||||||
continue
|
|
||||||
|
|
||||||
"""
|
|
||||||
if isinstance(model, UNet2DConditionModel):
|
|
||||||
if False and ".upsamplers." in m_name:
|
|
||||||
continue
|
continue
|
||||||
|
|
||||||
if False and ".downsamplers." in m_name:
|
if submodule_name == "conv2" and skip_conv2:
|
||||||
continue
|
continue
|
||||||
|
|
||||||
if True and ".resnets." in m_name:
|
m.asymmetric_padding_mode = {}
|
||||||
if True and ".conv1" in m_name:
|
m.asymmetric_padding = {}
|
||||||
if False and "down_blocks" in m_name:
|
m.asymmetric_padding_mode["x"] = "circular" if ("x" in seamless_axes) else "constant"
|
||||||
continue
|
m.asymmetric_padding["x"] = (
|
||||||
if False and "mid_block" in m_name:
|
m._reversed_padding_repeated_twice[0],
|
||||||
continue
|
m._reversed_padding_repeated_twice[1],
|
||||||
if False and "up_blocks" in m_name:
|
0,
|
||||||
continue
|
0,
|
||||||
|
)
|
||||||
|
m.asymmetric_padding_mode["y"] = "circular" if ("y" in seamless_axes) else "constant"
|
||||||
|
m.asymmetric_padding["y"] = (
|
||||||
|
0,
|
||||||
|
0,
|
||||||
|
m._reversed_padding_repeated_twice[2],
|
||||||
|
m._reversed_padding_repeated_twice[3],
|
||||||
|
)
|
||||||
|
|
||||||
if True and ".conv2" in m_name:
|
to_restore.append((m, m._conv_forward))
|
||||||
continue
|
m._conv_forward = _conv_forward_asymmetric.__get__(m, nn.Conv2d)
|
||||||
|
|
||||||
if True and ".conv_shortcut" in m_name:
|
|
||||||
continue
|
|
||||||
|
|
||||||
if True and ".attentions." in m_name:
|
|
||||||
continue
|
|
||||||
|
|
||||||
if False and m_name in ["conv_in", "conv_out"]:
|
|
||||||
continue
|
|
||||||
"""
|
|
||||||
|
|
||||||
if isinstance(m, (nn.Conv2d, nn.ConvTranspose2d)):
|
|
||||||
m.asymmetric_padding_mode = {}
|
|
||||||
m.asymmetric_padding = {}
|
|
||||||
m.asymmetric_padding_mode["x"] = "circular" if ("x" in seamless_axes) else "constant"
|
|
||||||
m.asymmetric_padding["x"] = (
|
|
||||||
m._reversed_padding_repeated_twice[0],
|
|
||||||
m._reversed_padding_repeated_twice[1],
|
|
||||||
0,
|
|
||||||
0,
|
|
||||||
)
|
|
||||||
m.asymmetric_padding_mode["y"] = "circular" if ("y" in seamless_axes) else "constant"
|
|
||||||
m.asymmetric_padding["y"] = (
|
|
||||||
0,
|
|
||||||
0,
|
|
||||||
m._reversed_padding_repeated_twice[2],
|
|
||||||
m._reversed_padding_repeated_twice[3],
|
|
||||||
)
|
|
||||||
|
|
||||||
to_restore.append((m, m._conv_forward))
|
|
||||||
m._conv_forward = _conv_forward_asymmetric.__get__(m, nn.Conv2d)
|
|
||||||
|
|
||||||
yield
|
yield
|
||||||
|
|
||||||
|
Loading…
Reference in New Issue
Block a user