diff --git a/invokeai/app/invocations/latent.py b/invokeai/app/invocations/latent.py index 80f90157df..c1918a981a 100644 --- a/invokeai/app/invocations/latent.py +++ b/invokeai/app/invocations/latent.py @@ -33,7 +33,7 @@ from invokeai.backend.model_management.models import ModelType, SilenceWarnings from ...backend.model_management.models import BaseModelType from ...backend.model_management.lora import ModelPatcher -from ...backend.model_management.seamless import set_unet_seamless, set_vae_seamless +from ...backend.model_management.seamless import set_seamless from ...backend.stable_diffusion import PipelineIntermediateState from ...backend.stable_diffusion.diffusers_pipeline import ( ConditioningData, @@ -401,7 +401,7 @@ class DenoiseLatentsInvocation(BaseInvocation): ) with ExitStack() as exit_stack, ModelPatcher.apply_lora_unet( unet_info.context.model, _lora_loader() - ), set_unet_seamless(unet_info.context.model, self.unet.seamless_axes), unet_info as unet: + ), set_seamless(unet_info.context.model, self.unet.seamless_axes), unet_info as unet: latents = latents.to(device=unet.device, dtype=unet.dtype) if noise is not None: noise = noise.to(device=unet.device, dtype=unet.dtype) @@ -491,7 +491,7 @@ class LatentsToImageInvocation(BaseInvocation): context=context, ) - with set_vae_seamless(vae_info.context.model, self.vae.seamless_axes), vae_info as vae: + with set_seamless(vae_info.context.model, self.vae.seamless_axes), vae_info as vae: latents = latents.to(vae.device) if self.fp32: vae.to(dtype=torch.float32) diff --git a/invokeai/app/invocations/model.py b/invokeai/app/invocations/model.py index c5bf4b1262..ca76ce7d51 100644 --- a/invokeai/app/invocations/model.py +++ b/invokeai/app/invocations/model.py @@ -398,8 +398,8 @@ class SeamlessModeOutput(BaseInvocationOutput): type: Literal["seamless_output"] = "seamless_output" # Outputs - unet: UNetField = OutputField(description=FieldDescriptions.unet, title="UNet") - vae: VaeField = OutputField(description=FieldDescriptions.vae, title="VAE") + unet: Optional[UNetField] = OutputField(description=FieldDescriptions.unet, title="UNet") + vae: Optional[VaeField] = OutputField(description=FieldDescriptions.vae, title="VAE") @title("Seamless") @tags("seamless", "model") diff --git a/invokeai/backend/model_management/seamless.py b/invokeai/backend/model_management/seamless.py index 49770b4281..db0274c3f7 100644 --- a/invokeai/backend/model_management/seamless.py +++ b/invokeai/backend/model_management/seamless.py @@ -1,7 +1,7 @@ from __future__ import annotations from contextlib import contextmanager - +from typing import TypeVar, Union import torch.nn as nn from diffusers.models import UNet2DModel, AutoencoderKL @@ -23,43 +23,10 @@ def _conv_forward_asymmetric(self, input, weight, bias): @contextmanager -def set_unet_seamless(model: UNet2DModel, seamless_axes): - try: - to_restore = [] - - for m in model.modules(): - 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) +ModelType = TypeVar('ModelType', UNet2DModel, AutoencoderKL) - yield - - finally: - for module, orig_conv_forward in to_restore: - module._conv_forward = orig_conv_forward - if hasattr(m, "asymmetric_padding_mode"): - del m.asymmetric_padding_mode - if hasattr(m, "asymmetric_padding"): - del m.asymmetric_padding - -def set_vae_seamless(model: AutoencoderKL, seamless_axes): +def set_seamless(model: ModelType, seamless_axes): try: to_restore = []