chore: import upstream snapshot with attribution

invokeai/backend/stable_diffusion/extensions/base.py

@@ -0,0 +1,72 @@
+from __future__ import annotations
+
+from contextlib import contextmanager
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, Callable, Dict, List
+
+from diffusers import UNet2DConditionModel
+
+if TYPE_CHECKING:
+    from invokeai.backend.stable_diffusion.denoise_context import DenoiseContext
+    from invokeai.backend.stable_diffusion.extension_callback_type import ExtensionCallbackType
+    from invokeai.backend.util.original_weights_storage import OriginalWeightsStorage
+
+
+@dataclass
+class CallbackMetadata:
+    callback_type: ExtensionCallbackType
+    order: int
+
+
+@dataclass
+class CallbackFunctionWithMetadata:
+    metadata: CallbackMetadata
+    function: Callable[[DenoiseContext], None]
+
+
+def callback(callback_type: ExtensionCallbackType, order: int = 0):
+    def _decorator(function):
+        function._ext_metadata = CallbackMetadata(
+            callback_type=callback_type,
+            order=order,
+        )
+        return function
+
+    return _decorator
+
+
+class ExtensionBase:
+    def __init__(self):
+        self._callbacks: Dict[ExtensionCallbackType, List[CallbackFunctionWithMetadata]] = {}
+
+        # Register all of the callback methods for this instance.
+        for func_name in dir(self):
+            func = getattr(self, func_name)
+            metadata = getattr(func, "_ext_metadata", None)
+            if metadata is not None and isinstance(metadata, CallbackMetadata):
+                if metadata.callback_type not in self._callbacks:
+                    self._callbacks[metadata.callback_type] = []
+                self._callbacks[metadata.callback_type].append(CallbackFunctionWithMetadata(metadata, func))
+
+    def get_callbacks(self):
+        return self._callbacks
+
+    @contextmanager
+    def patch_extension(self, ctx: DenoiseContext):
+        yield None
+
+    @contextmanager
+    def patch_unet(self, unet: UNet2DConditionModel, original_weights: OriginalWeightsStorage):
+        """A context manager for applying patches to the UNet model. The context manager's lifetime spans the entire
+        diffusion process. Weight unpatching is handled upstream, and is achieved by saving unchanged weights by
+        `original_weights.save` function. Note that this enables some performance optimization by avoiding redundant
+        operations. All other patches (e.g. changes to tensor shapes, function monkey-patches, etc.) should be unpatched
+        by this context manager.
+
+        Args:
+            unet (UNet2DConditionModel): The UNet model on execution device to patch.
+            original_weights (OriginalWeightsStorage): A storage with copy of the model's original weights in CPU, for
+                unpatching purposes. Extension should save tensor which being modified in this storage, also extensions
+                can access original weights values.
+        """
+        yield

invokeai/backend/stable_diffusion/extensions/controlnet.py

@@ -0,0 +1,158 @@
+from __future__ import annotations
+
+import math
+from contextlib import contextmanager
+from typing import TYPE_CHECKING, List, Optional, Union
+
+import torch
+from PIL.Image import Image
+
+from invokeai.app.invocations.constants import LATENT_SCALE_FACTOR
+from invokeai.app.util.controlnet_utils import CONTROLNET_MODE_VALUES, CONTROLNET_RESIZE_VALUES, prepare_control_image
+from invokeai.backend.stable_diffusion.denoise_context import UNetKwargs
+from invokeai.backend.stable_diffusion.diffusion.conditioning_data import ConditioningMode
+from invokeai.backend.stable_diffusion.extension_callback_type import ExtensionCallbackType
+from invokeai.backend.stable_diffusion.extensions.base import ExtensionBase, callback
+
+if TYPE_CHECKING:
+    from invokeai.backend.stable_diffusion.denoise_context import DenoiseContext
+    from invokeai.backend.util.hotfixes import ControlNetModel
+
+
+class ControlNetExt(ExtensionBase):
+    def __init__(
+        self,
+        model: ControlNetModel,
+        image: Image,
+        weight: Union[float, List[float]],
+        begin_step_percent: float,
+        end_step_percent: float,
+        control_mode: CONTROLNET_MODE_VALUES,
+        resize_mode: CONTROLNET_RESIZE_VALUES,
+    ):
+        super().__init__()
+        self._model = model
+        self._image = image
+        self._weight = weight
+        self._begin_step_percent = begin_step_percent
+        self._end_step_percent = end_step_percent
+        self._control_mode = control_mode
+        self._resize_mode = resize_mode
+
+        self._image_tensor: Optional[torch.Tensor] = None
+
+    @contextmanager
+    def patch_extension(self, ctx: DenoiseContext):
+        original_processors = self._model.attn_processors
+        try:
+            self._model.set_attn_processor(ctx.inputs.attention_processor_cls())
+
+            yield None
+        finally:
+            self._model.set_attn_processor(original_processors)
+
+    @callback(ExtensionCallbackType.PRE_DENOISE_LOOP)
+    def resize_image(self, ctx: DenoiseContext):
+        _, _, latent_height, latent_width = ctx.latents.shape
+        image_height = latent_height * LATENT_SCALE_FACTOR
+        image_width = latent_width * LATENT_SCALE_FACTOR
+
+        self._image_tensor = prepare_control_image(
+            image=self._image,
+            do_classifier_free_guidance=False,
+            width=image_width,
+            height=image_height,
+            device=ctx.latents.device,
+            dtype=ctx.latents.dtype,
+            control_mode=self._control_mode,
+            resize_mode=self._resize_mode,
+        )
+
+    @callback(ExtensionCallbackType.PRE_UNET)
+    def pre_unet_step(self, ctx: DenoiseContext):
+        # skip if model not active in current step
+        total_steps = len(ctx.inputs.timesteps)
+        first_step = math.floor(self._begin_step_percent * total_steps)
+        last_step = math.ceil(self._end_step_percent * total_steps)
+        if ctx.step_index < first_step or ctx.step_index > last_step:
+            return
+
+        # convert mode to internal flags
+        soft_injection = self._control_mode in ["more_prompt", "more_control"]
+        cfg_injection = self._control_mode in ["more_control", "unbalanced"]
+
+        # no negative conditioning in cfg_injection mode
+        if cfg_injection:
+            if ctx.conditioning_mode == ConditioningMode.Negative:
+                return
+            down_samples, mid_sample = self._run(ctx, soft_injection, ConditioningMode.Positive)
+
+            if ctx.conditioning_mode == ConditioningMode.Both:
+                # add zeros as samples for negative conditioning
+                down_samples = [torch.cat([torch.zeros_like(d), d]) for d in down_samples]
+                mid_sample = torch.cat([torch.zeros_like(mid_sample), mid_sample])
+
+        else:
+            down_samples, mid_sample = self._run(ctx, soft_injection, ctx.conditioning_mode)
+
+        if (
+            ctx.unet_kwargs.down_block_additional_residuals is None
+            and ctx.unet_kwargs.mid_block_additional_residual is None
+        ):
+            ctx.unet_kwargs.down_block_additional_residuals = down_samples
+            ctx.unet_kwargs.mid_block_additional_residual = mid_sample
+        else:
+            # add controlnet outputs together if have multiple controlnets
+            ctx.unet_kwargs.down_block_additional_residuals = [
+                samples_prev + samples_curr
+                for samples_prev, samples_curr in zip(
+                    ctx.unet_kwargs.down_block_additional_residuals, down_samples, strict=True
+                )
+            ]
+            ctx.unet_kwargs.mid_block_additional_residual += mid_sample
+
+    def _run(self, ctx: DenoiseContext, soft_injection: bool, conditioning_mode: ConditioningMode):
+        total_steps = len(ctx.inputs.timesteps)
+
+        model_input = ctx.latent_model_input
+        image_tensor = self._image_tensor
+        if conditioning_mode == ConditioningMode.Both:
+            model_input = torch.cat([model_input] * 2)
+            image_tensor = torch.cat([image_tensor] * 2)
+
+        cn_unet_kwargs = UNetKwargs(
+            sample=model_input,
+            timestep=ctx.timestep,
+            encoder_hidden_states=None,  # set later by conditioning
+            cross_attention_kwargs=dict(  # noqa: C408
+                percent_through=ctx.step_index / total_steps,
+            ),
+        )
+
+        ctx.inputs.conditioning_data.to_unet_kwargs(cn_unet_kwargs, conditioning_mode=conditioning_mode)
+
+        # get static weight, or weight corresponding to current step
+        weight = self._weight
+        if isinstance(weight, list):
+            weight = weight[ctx.step_index]
+
+        tmp_kwargs = vars(cn_unet_kwargs)
+
+        # Remove kwargs not related to ControlNet unet
+        # ControlNet guidance fields
+        del tmp_kwargs["down_block_additional_residuals"]
+        del tmp_kwargs["mid_block_additional_residual"]
+
+        # T2i Adapter guidance fields
+        del tmp_kwargs["down_intrablock_additional_residuals"]
+
+        # controlnet(s) inference
+        down_samples, mid_sample = self._model(
+            controlnet_cond=image_tensor,
+            conditioning_scale=weight,  # controlnet specific, NOT the guidance scale
+            guess_mode=soft_injection,  # this is still called guess_mode in diffusers ControlNetModel
+            return_dict=False,
+            **vars(cn_unet_kwargs),
+        )
+
+        return down_samples, mid_sample

invokeai/backend/stable_diffusion/extensions/freeu.py

@@ -0,0 +1,35 @@
+from __future__ import annotations
+
+from contextlib import contextmanager
+from typing import TYPE_CHECKING
+
+from diffusers import UNet2DConditionModel
+
+from invokeai.backend.stable_diffusion.extensions.base import ExtensionBase
+
+if TYPE_CHECKING:
+    from invokeai.app.shared.models import FreeUConfig
+    from invokeai.backend.util.original_weights_storage import OriginalWeightsStorage
+
+
+class FreeUExt(ExtensionBase):
+    def __init__(
+        self,
+        freeu_config: FreeUConfig,
+    ):
+        super().__init__()
+        self._freeu_config = freeu_config
+
+    @contextmanager
+    def patch_unet(self, unet: UNet2DConditionModel, original_weights: OriginalWeightsStorage):
+        unet.enable_freeu(
+            b1=self._freeu_config.b1,
+            b2=self._freeu_config.b2,
+            s1=self._freeu_config.s1,
+            s2=self._freeu_config.s2,
+        )
+
+        try:
+            yield
+        finally:
+            unet.disable_freeu()

invokeai/backend/stable_diffusion/extensions/inpaint.py

@@ -0,0 +1,120 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Optional
+
+import einops
+import torch
+from diffusers import UNet2DConditionModel
+
+from invokeai.backend.stable_diffusion.extension_callback_type import ExtensionCallbackType
+from invokeai.backend.stable_diffusion.extensions.base import ExtensionBase, callback
+
+if TYPE_CHECKING:
+    from invokeai.backend.stable_diffusion.denoise_context import DenoiseContext
+
+
+class InpaintExt(ExtensionBase):
+    """An extension for inpainting with non-inpainting models. See `InpaintModelExt` for inpainting with inpainting
+    models.
+    """
+
+    def __init__(
+        self,
+        mask: torch.Tensor,
+        is_gradient_mask: bool,
+    ):
+        """Initialize InpaintExt.
+        Args:
+            mask (torch.Tensor): The inpainting mask. Shape: (1, 1, latent_height, latent_width). Values are
+                expected to be in the range [0, 1]. A value of 1 means that the corresponding 'pixel' should not be
+                inpainted.
+            is_gradient_mask (bool): If True, mask is interpreted as a gradient mask meaning that the mask values range
+                from 0 to 1. If False, mask is interpreted as binary mask meaning that the mask values are either 0 or
+                1.
+        """
+        super().__init__()
+        self._mask = mask
+        self._is_gradient_mask = is_gradient_mask
+
+        # Noise, which used to noisify unmasked part of image
+        # if noise provided to context, then it will be used
+        # if no noise provided, then noise will be generated based on seed
+        self._noise: Optional[torch.Tensor] = None
+
+    @staticmethod
+    def _is_normal_model(unet: UNet2DConditionModel):
+        """Checks if the provided UNet belongs to a regular model.
+        The `in_channels` of a UNet vary depending on model type:
+        - normal - 4
+        - depth - 5
+        - inpaint - 9
+        """
+        return unet.conv_in.in_channels == 4
+
+    def _apply_mask(self, ctx: DenoiseContext, latents: torch.Tensor, t: torch.Tensor) -> torch.Tensor:
+        batch_size = latents.size(0)
+        mask = einops.repeat(self._mask, "b c h w -> (repeat b) c h w", repeat=batch_size)
+        if t.dim() == 0:
+            # some schedulers expect t to be one-dimensional.
+            # TODO: file diffusers bug about inconsistency?
+            t = einops.repeat(t, "-> batch", batch=batch_size)
+        # Noise shouldn't be re-randomized between steps here. The multistep schedulers
+        # get very confused about what is happening from step to step when we do that.
+        mask_latents = ctx.scheduler.add_noise(ctx.inputs.orig_latents, self._noise, t)
+        # TODO: Do we need to also apply scheduler.scale_model_input? Or is add_noise appropriately scaled already?
+        # mask_latents = self.scheduler.scale_model_input(mask_latents, t)
+        mask_latents = einops.repeat(mask_latents, "b c h w -> (repeat b) c h w", repeat=batch_size)
+        if self._is_gradient_mask:
+            threshold = (t.item()) / ctx.scheduler.config.num_train_timesteps
+            mask_bool = mask < 1 - threshold
+            masked_input = torch.where(mask_bool, latents, mask_latents)
+        else:
+            masked_input = torch.lerp(latents, mask_latents.to(dtype=latents.dtype), mask.to(dtype=latents.dtype))
+        return masked_input
+
+    @callback(ExtensionCallbackType.PRE_DENOISE_LOOP)
+    def init_tensors(self, ctx: DenoiseContext):
+        if not self._is_normal_model(ctx.unet):
+            raise ValueError(
+                "InpaintExt should be used only on normal (non-inpainting) models. This could be caused by an "
+                "inpainting model that was incorrectly marked as a non-inpainting model. In some cases, this can be "
+                "fixed by removing and re-adding the model (so that it gets re-probed)."
+            )
+
+        self._mask = self._mask.to(device=ctx.latents.device, dtype=ctx.latents.dtype)
+
+        self._noise = ctx.inputs.noise
+        # 'noise' might be None if the latents have already been noised (e.g. when running the SDXL refiner).
+        # We still need noise for inpainting, so we generate it from the seed here.
+        if self._noise is None:
+            self._noise = torch.randn(
+                ctx.latents.shape,
+                dtype=torch.float32,
+                device="cpu",
+                generator=torch.Generator(device="cpu").manual_seed(ctx.seed),
+            ).to(device=ctx.latents.device, dtype=ctx.latents.dtype)
+
+    # Use negative order to make extensions with default order work with patched latents
+    @callback(ExtensionCallbackType.PRE_STEP, order=-100)
+    def apply_mask_to_initial_latents(self, ctx: DenoiseContext):
+        ctx.latents = self._apply_mask(ctx, ctx.latents, ctx.timestep)
+
+    # TODO: redo this with preview events rewrite
+    # Use negative order to make extensions with default order work with patched latents
+    @callback(ExtensionCallbackType.POST_STEP, order=-100)
+    def apply_mask_to_step_output(self, ctx: DenoiseContext):
+        timestep = ctx.scheduler.timesteps[-1]
+        if hasattr(ctx.step_output, "denoised"):
+            ctx.step_output.denoised = self._apply_mask(ctx, ctx.step_output.denoised, timestep)
+        elif hasattr(ctx.step_output, "pred_original_sample"):
+            ctx.step_output.pred_original_sample = self._apply_mask(ctx, ctx.step_output.pred_original_sample, timestep)
+        else:
+            ctx.step_output.pred_original_sample = self._apply_mask(ctx, ctx.step_output.prev_sample, timestep)
+
+    # Restore unmasked part after the last step is completed
+    @callback(ExtensionCallbackType.POST_DENOISE_LOOP)
+    def restore_unmasked(self, ctx: DenoiseContext):
+        if self._is_gradient_mask:
+            ctx.latents = torch.where(self._mask < 1, ctx.latents, ctx.inputs.orig_latents)
+        else:
+            ctx.latents = torch.lerp(ctx.latents, ctx.inputs.orig_latents, self._mask)

invokeai/backend/stable_diffusion/extensions/inpaint_model.py

@@ -0,0 +1,88 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Optional
+
+import torch
+from diffusers import UNet2DConditionModel
+
+from invokeai.backend.stable_diffusion.extension_callback_type import ExtensionCallbackType
+from invokeai.backend.stable_diffusion.extensions.base import ExtensionBase, callback
+
+if TYPE_CHECKING:
+    from invokeai.backend.stable_diffusion.denoise_context import DenoiseContext
+
+
+class InpaintModelExt(ExtensionBase):
+    """An extension for inpainting with inpainting models. See `InpaintExt` for inpainting with non-inpainting
+    models.
+    """
+
+    def __init__(
+        self,
+        mask: Optional[torch.Tensor],
+        masked_latents: Optional[torch.Tensor],
+        is_gradient_mask: bool,
+    ):
+        """Initialize InpaintModelExt.
+        Args:
+            mask (Optional[torch.Tensor]): The inpainting mask. Shape: (1, 1, latent_height, latent_width). Values are
+                expected to be in the range [0, 1]. A value of 1 means that the corresponding 'pixel' should not be
+                inpainted.
+            masked_latents (Optional[torch.Tensor]): Latents of initial image, with masked out by black color inpainted area.
+                If mask provided, then too should be provided. Shape: (1, 1, latent_height, latent_width)
+            is_gradient_mask (bool): If True, mask is interpreted as a gradient mask meaning that the mask values range
+                from 0 to 1. If False, mask is interpreted as binary mask meaning that the mask values are either 0 or
+                1.
+        """
+        super().__init__()
+        if mask is not None and masked_latents is None:
+            raise ValueError("Source image required for inpaint mask when inpaint model used!")
+
+        # Inverse mask, because inpaint models treat mask as: 0 - remain same, 1 - inpaint
+        self._mask = None
+        if mask is not None:
+            self._mask = 1 - mask
+        self._masked_latents = masked_latents
+        self._is_gradient_mask = is_gradient_mask
+
+    @staticmethod
+    def _is_inpaint_model(unet: UNet2DConditionModel):
+        """Checks if the provided UNet belongs to a regular model.
+        The `in_channels` of a UNet vary depending on model type:
+        - normal - 4
+        - depth - 5
+        - inpaint - 9
+        """
+        return unet.conv_in.in_channels == 9
+
+    @callback(ExtensionCallbackType.PRE_DENOISE_LOOP)
+    def init_tensors(self, ctx: DenoiseContext):
+        if not self._is_inpaint_model(ctx.unet):
+            raise ValueError("InpaintModelExt should be used only on inpaint models!")
+
+        if self._mask is None:
+            self._mask = torch.ones_like(ctx.latents[:1, :1])
+        self._mask = self._mask.to(device=ctx.latents.device, dtype=ctx.latents.dtype)
+
+        if self._masked_latents is None:
+            self._masked_latents = torch.zeros_like(ctx.latents[:1])
+        self._masked_latents = self._masked_latents.to(device=ctx.latents.device, dtype=ctx.latents.dtype)
+
+    # Do last so that other extensions works with normal latents
+    @callback(ExtensionCallbackType.PRE_UNET, order=1000)
+    def append_inpaint_layers(self, ctx: DenoiseContext):
+        batch_size = ctx.unet_kwargs.sample.shape[0]
+        b_mask = torch.cat([self._mask] * batch_size)
+        b_masked_latents = torch.cat([self._masked_latents] * batch_size)
+        ctx.unet_kwargs.sample = torch.cat(
+            [ctx.unet_kwargs.sample, b_mask, b_masked_latents],
+            dim=1,
+        )
+
+    # Restore unmasked part as inpaint model can change unmasked part slightly
+    @callback(ExtensionCallbackType.POST_DENOISE_LOOP)
+    def restore_unmasked(self, ctx: DenoiseContext):
+        if self._is_gradient_mask:
+            ctx.latents = torch.where(self._mask > 0, ctx.latents, ctx.inputs.orig_latents)
+        else:
+            ctx.latents = torch.lerp(ctx.inputs.orig_latents, ctx.latents, self._mask)

invokeai/backend/stable_diffusion/extensions/lora.py

@@ -0,0 +1,47 @@
+from __future__ import annotations
+
+from contextlib import contextmanager
+from typing import TYPE_CHECKING
+
+from diffusers import UNet2DConditionModel
+
+from invokeai.backend.patches.layer_patcher import LayerPatcher
+from invokeai.backend.patches.model_patch_raw import ModelPatchRaw
+from invokeai.backend.stable_diffusion.extensions.base import ExtensionBase
+
+if TYPE_CHECKING:
+    from invokeai.app.invocations.model import ModelIdentifierField
+    from invokeai.app.services.shared.invocation_context import InvocationContext
+    from invokeai.backend.util.original_weights_storage import OriginalWeightsStorage
+
+
+class LoRAExt(ExtensionBase):
+    def __init__(
+        self,
+        node_context: InvocationContext,
+        model_id: ModelIdentifierField,
+        weight: float,
+    ):
+        super().__init__()
+        self._node_context = node_context
+        self._model_id = model_id
+        self._weight = weight
+
+    @contextmanager
+    def patch_unet(self, unet: UNet2DConditionModel, original_weights: OriginalWeightsStorage):
+        lora_model = self._node_context.models.load(self._model_id).model
+        assert isinstance(lora_model, ModelPatchRaw)
+        LayerPatcher.apply_smart_model_patch(
+            model=unet,
+            prefix="lora_unet_",
+            patch=lora_model,
+            patch_weight=self._weight,
+            original_weights=original_weights,
+            original_modules={},
+            dtype=unet.dtype,
+            force_direct_patching=True,
+            force_sidecar_patching=False,
+        )
+        del lora_model
+
+        yield

invokeai/backend/stable_diffusion/extensions/preview.py

@@ -0,0 +1,63 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, Callable, Optional
+
+import torch
+
+from invokeai.backend.stable_diffusion.extension_callback_type import ExtensionCallbackType
+from invokeai.backend.stable_diffusion.extensions.base import ExtensionBase, callback
+
+if TYPE_CHECKING:
+    from invokeai.backend.stable_diffusion.denoise_context import DenoiseContext
+
+
+# TODO: change event to accept image instead of latents
+@dataclass
+class PipelineIntermediateState:
+    step: int
+    order: int
+    total_steps: int
+    timestep: int
+    latents: torch.Tensor
+    predicted_original: Optional[torch.Tensor] = None
+
+
+class PreviewExt(ExtensionBase):
+    def __init__(self, callback: Callable[[PipelineIntermediateState], None]):
+        super().__init__()
+        self.callback = callback
+
+    # do last so that all other changes shown
+    @callback(ExtensionCallbackType.PRE_DENOISE_LOOP, order=1000)
+    def initial_preview(self, ctx: DenoiseContext):
+        self.callback(
+            PipelineIntermediateState(
+                step=0,
+                order=ctx.scheduler.order,
+                total_steps=len(ctx.inputs.timesteps),
+                timestep=int(ctx.scheduler.config.num_train_timesteps),  # TODO: is there any code which uses it?
+                latents=ctx.latents,
+            )
+        )
+
+    # do last so that all other changes shown
+    @callback(ExtensionCallbackType.POST_STEP, order=1000)
+    def step_preview(self, ctx: DenoiseContext):
+        if hasattr(ctx.step_output, "denoised"):
+            predicted_original = ctx.step_output.denoised
+        elif hasattr(ctx.step_output, "pred_original_sample"):
+            predicted_original = ctx.step_output.pred_original_sample
+        else:
+            predicted_original = ctx.step_output.prev_sample
+
+        self.callback(
+            PipelineIntermediateState(
+                step=ctx.step_index,
+                order=ctx.scheduler.order,
+                total_steps=len(ctx.inputs.timesteps),
+                timestep=int(ctx.timestep),  # TODO: is there any code which uses it?
+                latents=ctx.step_output.prev_sample,
+                predicted_original=predicted_original,  # TODO: is there any reason for additional field?
+            )
+        )

invokeai/backend/stable_diffusion/extensions/rescale_cfg.py

@@ -0,0 +1,36 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+import torch
+
+from invokeai.backend.stable_diffusion.extension_callback_type import ExtensionCallbackType
+from invokeai.backend.stable_diffusion.extensions.base import ExtensionBase, callback
+
+if TYPE_CHECKING:
+    from invokeai.backend.stable_diffusion.denoise_context import DenoiseContext
+
+
+class RescaleCFGExt(ExtensionBase):
+    def __init__(self, rescale_multiplier: float):
+        super().__init__()
+        self._rescale_multiplier = rescale_multiplier
+
+    @staticmethod
+    def _rescale_cfg(total_noise_pred: torch.Tensor, pos_noise_pred: torch.Tensor, multiplier: float = 0.7):
+        """Implementation of Algorithm 2 from https://arxiv.org/pdf/2305.08891.pdf."""
+        ro_pos = torch.std(pos_noise_pred, dim=(1, 2, 3), keepdim=True)
+        ro_cfg = torch.std(total_noise_pred, dim=(1, 2, 3), keepdim=True)
+
+        x_rescaled = total_noise_pred * (ro_pos / ro_cfg)
+        x_final = multiplier * x_rescaled + (1.0 - multiplier) * total_noise_pred
+        return x_final
+
+    @callback(ExtensionCallbackType.POST_COMBINE_NOISE_PREDS)
+    def rescale_noise_pred(self, ctx: DenoiseContext):
+        if self._rescale_multiplier > 0:
+            ctx.noise_pred = self._rescale_cfg(
+                ctx.noise_pred,
+                ctx.positive_noise_pred,
+                self._rescale_multiplier,
+            )

invokeai/backend/stable_diffusion/extensions/seamless.py

@@ -0,0 +1,71 @@
+from __future__ import annotations
+
+from contextlib import contextmanager
+from typing import Callable, Dict, List, Optional, Tuple
+
+import torch
+import torch.nn as nn
+from diffusers import UNet2DConditionModel
+from diffusers.models.lora import LoRACompatibleConv
+
+from invokeai.backend.stable_diffusion.extensions.base import ExtensionBase
+
+
+class SeamlessExt(ExtensionBase):
+    def __init__(
+        self,
+        seamless_axes: List[str],
+    ):
+        super().__init__()
+        self._seamless_axes = seamless_axes
+
+    @contextmanager
+    def patch_unet(self, unet: UNet2DConditionModel, cached_weights: Optional[Dict[str, torch.Tensor]] = None):
+        with self.static_patch_model(
+            model=unet,
+            seamless_axes=self._seamless_axes,
+        ):
+            yield
+
+    @staticmethod
+    @contextmanager
+    def static_patch_model(
+        model: torch.nn.Module,
+        seamless_axes: List[str],
+    ):
+        if not seamless_axes:
+            yield
+            return
+
+        x_mode = "circular" if "x" in seamless_axes else "constant"
+        y_mode = "circular" if "y" in seamless_axes else "constant"
+
+        # override conv_forward
+        # https://github.com/huggingface/diffusers/issues/556#issuecomment-1993287019
+        def _conv_forward_asymmetric(
+            self, input: torch.Tensor, weight: torch.Tensor, bias: Optional[torch.Tensor] = None
+        ):
+            self.paddingX = (self._reversed_padding_repeated_twice[0], self._reversed_padding_repeated_twice[1], 0, 0)
+            self.paddingY = (0, 0, self._reversed_padding_repeated_twice[2], self._reversed_padding_repeated_twice[3])
+            working = torch.nn.functional.pad(input, self.paddingX, mode=x_mode)
+            working = torch.nn.functional.pad(working, self.paddingY, mode=y_mode)
+            return torch.nn.functional.conv2d(
+                working, weight, bias, self.stride, torch.nn.modules.utils._pair(0), self.dilation, self.groups
+            )
+
+        original_layers: List[Tuple[nn.Conv2d, Callable]] = []
+        try:
+            for layer in model.modules():
+                if not isinstance(layer, torch.nn.Conv2d):
+                    continue
+
+                if isinstance(layer, LoRACompatibleConv) and layer.lora_layer is None:
+                    layer.lora_layer = lambda *x: 0
+                original_layers.append((layer, layer._conv_forward))
+                layer._conv_forward = _conv_forward_asymmetric.__get__(layer, torch.nn.Conv2d)
+
+            yield
+
+        finally:
+            for layer, orig_conv_forward in original_layers:
+                layer._conv_forward = orig_conv_forward

invokeai/backend/stable_diffusion/extensions/t2i_adapter.py

@@ -0,0 +1,121 @@
+from __future__ import annotations
+
+import math
+from typing import TYPE_CHECKING, List, Optional, Union
+
+import torch
+from diffusers import T2IAdapter
+from PIL.Image import Image
+
+from invokeai.app.util.controlnet_utils import prepare_control_image
+from invokeai.backend.model_manager.taxonomy import BaseModelType
+from invokeai.backend.stable_diffusion.diffusion.conditioning_data import ConditioningMode
+from invokeai.backend.stable_diffusion.extension_callback_type import ExtensionCallbackType
+from invokeai.backend.stable_diffusion.extensions.base import ExtensionBase, callback
+from invokeai.backend.util.devices import TorchDevice
+
+if TYPE_CHECKING:
+    from invokeai.app.invocations.model import ModelIdentifierField
+    from invokeai.app.services.shared.invocation_context import InvocationContext
+    from invokeai.app.util.controlnet_utils import CONTROLNET_RESIZE_VALUES
+    from invokeai.backend.stable_diffusion.denoise_context import DenoiseContext
+
+
+class T2IAdapterExt(ExtensionBase):
+    def __init__(
+        self,
+        node_context: InvocationContext,
+        model_id: ModelIdentifierField,
+        image: Image,
+        weight: Union[float, List[float]],
+        begin_step_percent: float,
+        end_step_percent: float,
+        resize_mode: CONTROLNET_RESIZE_VALUES,
+    ):
+        super().__init__()
+        self._node_context = node_context
+        self._model_id = model_id
+        self._image = image
+        self._weight = weight
+        self._resize_mode = resize_mode
+        self._begin_step_percent = begin_step_percent
+        self._end_step_percent = end_step_percent
+
+        self._adapter_state: Optional[List[torch.Tensor]] = None
+
+        # The max_unet_downscale is the maximum amount that the UNet model downscales the latent image internally.
+        model_config = self._node_context.models.get_config(self._model_id.key)
+        if model_config.base == BaseModelType.StableDiffusion1:
+            self._max_unet_downscale = 8
+        elif model_config.base == BaseModelType.StableDiffusionXL:
+            self._max_unet_downscale = 4
+        else:
+            raise ValueError(f"Unexpected T2I-Adapter base model type: '{model_config.base}'.")
+
+    @callback(ExtensionCallbackType.SETUP)
+    def setup(self, ctx: DenoiseContext):
+        t2i_model: T2IAdapter
+        with self._node_context.models.load(self._model_id) as t2i_model:
+            _, _, latents_height, latents_width = ctx.inputs.orig_latents.shape
+
+            self._adapter_state = self._run_model(
+                model=t2i_model,
+                image=self._image,
+                latents_height=latents_height,
+                latents_width=latents_width,
+            )
+
+    def _run_model(
+        self,
+        model: T2IAdapter,
+        image: Image,
+        latents_height: int,
+        latents_width: int,
+    ):
+        # Resize the T2I-Adapter input image.
+        # We select the resize dimensions so that after the T2I-Adapter's total_downscale_factor is applied, the
+        # result will match the latent image's dimensions after max_unet_downscale is applied.
+        input_height = latents_height // self._max_unet_downscale * model.total_downscale_factor
+        input_width = latents_width // self._max_unet_downscale * model.total_downscale_factor
+
+        # Note: We have hard-coded `do_classifier_free_guidance=False`. This is because we only want to prepare
+        # a single image. If CFG is enabled, we will duplicate the resultant tensor after applying the
+        # T2I-Adapter model.
+        #
+        # Note: We re-use the `prepare_control_image(...)` from ControlNet for T2I-Adapter, because it has many
+        # of the same requirements (e.g. preserving binary masks during resize).
+        t2i_image = prepare_control_image(
+            image=image,
+            do_classifier_free_guidance=False,
+            width=input_width,
+            height=input_height,
+            num_channels=model.config["in_channels"],
+            device=TorchDevice.choose_torch_device(),
+            dtype=model.dtype,
+            resize_mode=self._resize_mode,
+        )
+
+        return model(t2i_image)
+
+    @callback(ExtensionCallbackType.PRE_UNET)
+    def pre_unet_step(self, ctx: DenoiseContext):
+        # skip if model not active in current step
+        total_steps = len(ctx.inputs.timesteps)
+        first_step = math.floor(self._begin_step_percent * total_steps)
+        last_step = math.ceil(self._end_step_percent * total_steps)
+        if ctx.step_index < first_step or ctx.step_index > last_step:
+            return
+
+        weight = self._weight
+        if isinstance(weight, list):
+            weight = weight[ctx.step_index]
+
+        adapter_state = self._adapter_state
+        if ctx.conditioning_mode == ConditioningMode.Both:
+            adapter_state = [torch.cat([v] * 2) for v in adapter_state]
+
+        if ctx.unet_kwargs.down_intrablock_additional_residuals is None:
+            ctx.unet_kwargs.down_intrablock_additional_residuals = [v * weight for v in adapter_state]
+        else:
+            for i, value in enumerate(adapter_state):
+                ctx.unet_kwargs.down_intrablock_additional_residuals[i] += value * weight