Ryan's suggested changes to extension manager/extensions

Co-Authored-By: Ryan Dick <14897797+RyanJDick@users.noreply.github.com>

invokeai/app/invocations/denoise_latents.py

@@ -57,6 +57,7 @@ from invokeai.backend.stable_diffusion.diffusion.conditioning_data import (
 )
 from invokeai.backend.stable_diffusion.diffusion.custom_atttention import CustomAttnProcessor2_0
 from invokeai.backend.stable_diffusion.diffusion_backend import StableDiffusionBackend
+from invokeai.backend.stable_diffusion.extension_callback_type import ExtensionCallbackType
 from invokeai.backend.stable_diffusion.extensions.preview import PreviewExt
 from invokeai.backend.stable_diffusion.extensions_manager import ExtensionsManager
 from invokeai.backend.stable_diffusion.schedulers import SCHEDULER_MAP
@@ -790,7 +791,7 @@ class DenoiseLatentsInvocation(BaseInvocation):
         ext_manager.add_extension(PreviewExt(step_callback))
 
         # ext: t2i/ip adapter
-        ext_manager.callbacks.setup(denoise_ctx, ext_manager)
+        ext_manager.run_callback(ExtensionCallbackType.SETUP, denoise_ctx)
 
         unet_info = context.models.load(self.unet.unet)
         assert isinstance(unet_info.model, UNet2DConditionModel)

invokeai/backend/stable_diffusion/diffusion_backend.py

@@ -8,6 +8,7 @@ from tqdm.auto import tqdm
 from invokeai.app.services.config.config_default import get_config
 from invokeai.backend.stable_diffusion.denoise_context import DenoiseContext, 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_manager import ExtensionsManager
 
 
@@ -41,23 +42,23 @@ class StableDiffusionBackend:
 
         # ext: inpaint[pre_denoise_loop, priority=normal] (maybe init, but not sure if it needed)
         # ext: preview[pre_denoise_loop, priority=low]
-        ext_manager.callbacks.pre_denoise_loop(ctx, ext_manager)
+        ext_manager.run_callback(ExtensionCallbackType.PRE_DENOISE_LOOP, ctx)
 
         for ctx.step_index, ctx.timestep in enumerate(tqdm(ctx.inputs.timesteps)):  # noqa: B020
             # ext: inpaint (apply mask to latents on non-inpaint models)
-            ext_manager.callbacks.pre_step(ctx, ext_manager)
+            ext_manager.run_callback(ExtensionCallbackType.PRE_STEP, ctx)
 
             # ext: tiles? [override: step]
             ctx.step_output = self.step(ctx, ext_manager)
 
             # ext: inpaint[post_step, priority=high] (apply mask to preview on non-inpaint models)
             # ext: preview[post_step, priority=low]
-            ext_manager.callbacks.post_step(ctx, ext_manager)
+            ext_manager.run_callback(ExtensionCallbackType.POST_STEP, ctx)
 
             ctx.latents = ctx.step_output.prev_sample
 
         # ext: inpaint[post_denoise_loop] (restore unmasked part)
-        ext_manager.callbacks.post_denoise_loop(ctx, ext_manager)
+        ext_manager.run_callback(ExtensionCallbackType.POST_DENOISE_LOOP, ctx)
         return ctx.latents
 
     @torch.inference_mode()
@@ -80,7 +81,7 @@ class StableDiffusionBackend:
 
         # ext: cfg_rescale [modify_noise_prediction]
         # TODO: rename
-        ext_manager.callbacks.post_apply_cfg(ctx, ext_manager)
+        ext_manager.run_callback(ExtensionCallbackType.POST_APPLY_CFG, ctx)
 
         # compute the previous noisy sample x_t -> x_t-1
         step_output = ctx.scheduler.step(ctx.noise_pred, ctx.timestep, ctx.latents, **ctx.inputs.scheduler_step_kwargs)
@@ -120,14 +121,14 @@ class StableDiffusionBackend:
         ctx.inputs.conditioning_data.to_unet_kwargs(ctx.unet_kwargs, ctx.conditioning_mode)
 
         # ext: controlnet/ip/t2i [pre_unet]
-        ext_manager.callbacks.pre_unet(ctx, ext_manager)
+        ext_manager.run_callback(ExtensionCallbackType.PRE_UNET, ctx)
 
         # ext: inpaint [pre_unet, priority=low]
         # or
         # ext: inpaint [override: unet_forward]
         noise_pred = self._unet_forward(**vars(ctx.unet_kwargs))
 
-        ext_manager.callbacks.post_unet(ctx, ext_manager)
+        ext_manager.run_callback(ExtensionCallbackType.POST_UNET, ctx)
 
         # clean up locals
         ctx.unet_kwargs = None

invokeai/backend/stable_diffusion/extension_callback_type.py

@@ -0,0 +1,12 @@
+from enum import Enum
+
+
+class ExtensionCallbackType(Enum):
+    SETUP = "setup"
+    PRE_DENOISE_LOOP = "pre_denoise_loop"
+    POST_DENOISE_LOOP = "post_denoise_loop"
+    PRE_STEP = "pre_step"
+    POST_STEP = "post_step"
+    PRE_UNET = "pre_unet"
+    POST_UNET = "post_unet"
+    POST_APPLY_CFG = "post_apply_cfg"

invokeai/backend/stable_diffusion/extensions/base.py

@@ -2,44 +2,54 @@ from __future__ import annotations
 
 from contextlib import contextmanager
 from dataclasses import dataclass
-from typing import TYPE_CHECKING, Callable, Dict, List, Optional
+from typing import TYPE_CHECKING, Callable, Dict, List
 
 import torch
 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
 
 
 @dataclass
-class InjectionInfo:
-    type: str
-    name: str
-    order: Optional[int]
-    function: Callable
+class CallbackMetadata:
+    callback_type: ExtensionCallbackType
+    order: int
 
 
-def callback(name: str, order: int = 0):
-    def _decorator(func):
-        func.__inj_info__ = {
-            "type": "callback",
-            "name": name,
-            "order": order,
-        }
-        return func
+@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.injections: List[InjectionInfo] = []
+        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)
-            if not callable(func) or not hasattr(func, "__inj_info__"):
-                continue
+            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))
 
-            self.injections.append(InjectionInfo(**func.__inj_info__, function=func))
+    def get_callbacks(self):
+        return self._callbacks
 
     @contextmanager
     def patch_extension(self, context: DenoiseContext):

invokeai/backend/stable_diffusion/extensions/preview.py

@@ -5,11 +5,11 @@ 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
-    from invokeai.backend.stable_diffusion.extensions_manager import ExtensionsManager
 
 
 # TODO: change event to accept image instead of latents
@@ -29,8 +29,8 @@ class PreviewExt(ExtensionBase):
         self.callback = callback
 
     # do last so that all other changes shown
-    @callback("pre_denoise_loop", order=1000)
-    def initial_preview(self, ctx: DenoiseContext, ext_manager: ExtensionsManager):
+    @callback(ExtensionCallbackType.PRE_DENOISE_LOOP, order=1000)
+    def initial_preview(self, ctx: DenoiseContext):
         self.callback(
             PipelineIntermediateState(
                 step=-1,
@@ -42,8 +42,8 @@ class PreviewExt(ExtensionBase):
         )
 
     # do last so that all other changes shown
-    @callback("post_step", order=1000)
-    def step_preview(self, ctx: DenoiseContext, ext_manager: ExtensionsManager):
+    @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"):

invokeai/backend/stable_diffusion/extensions_manager.py

@@ -1,8 +1,6 @@
 from __future__ import annotations
 
-from abc import ABC, abstractmethod
 from contextlib import ExitStack, contextmanager
-from functools import partial
 from typing import TYPE_CHECKING, Callable, Dict, List, Optional
 
 import torch
@@ -12,102 +10,50 @@ from invokeai.app.services.session_processor.session_processor_common import Can
 
 if TYPE_CHECKING:
     from invokeai.backend.stable_diffusion.denoise_context import DenoiseContext
-    from invokeai.backend.stable_diffusion.extensions.base import ExtensionBase
-
-
-class ExtCallbacksApi(ABC):
-    @abstractmethod
-    def setup(self, ctx: DenoiseContext, ext_manager: ExtensionsManager):
-        pass
-
-    @abstractmethod
-    def pre_denoise_loop(self, ctx: DenoiseContext, ext_manager: ExtensionsManager):
-        pass
-
-    @abstractmethod
-    def post_denoise_loop(self, ctx: DenoiseContext, ext_manager: ExtensionsManager):
-        pass
-
-    @abstractmethod
-    def pre_step(self, ctx: DenoiseContext, ext_manager: ExtensionsManager):
-        pass
-
-    @abstractmethod
-    def post_step(self, ctx: DenoiseContext, ext_manager: ExtensionsManager):
-        pass
-
-    @abstractmethod
-    def pre_unet(self, ctx: DenoiseContext, ext_manager: ExtensionsManager):
-        pass
-
-    @abstractmethod
-    def post_unet(self, ctx: DenoiseContext, ext_manager: ExtensionsManager):
-        pass
-
-    @abstractmethod
-    def post_apply_cfg(self, ctx: DenoiseContext, ext_manager: ExtensionsManager):
-        pass
-
-
-class ProxyCallsClass:
-    def __init__(self, handler):
-        self._handler = handler
-
-    def __getattr__(self, item):
-        return partial(self._handler, item)
-
-
-class CallbackInjectionPoint:
-    def __init__(self):
-        self.handlers = {}
-
-    def add(self, func: Callable, order: int):
-        if order not in self.handlers:
-            self.handlers[order] = []
-        self.handlers[order].append(func)
-
-    def __call__(self, *args, **kwargs):
-        for order in sorted(self.handlers.keys(), reverse=True):
-            for handler in self.handlers[order]:
-                handler(*args, **kwargs)
+    from invokeai.backend.stable_diffusion.extension_callback_type import ExtensionCallbackType
+    from invokeai.backend.stable_diffusion.extensions.base import CallbackFunctionWithMetadata, ExtensionBase
 
 
 class ExtensionsManager:
     def __init__(self, is_canceled: Optional[Callable[[], bool]] = None):
-        self.extensions: List[ExtensionBase] = []
         self._is_canceled = is_canceled
 
-        self._callbacks: Dict[str, CallbackInjectionPoint] = {}
-        self.callbacks: ExtCallbacksApi = ProxyCallsClass(self.call_callback)
+        self._extensions: List[ExtensionBase] = []
+        self._ordered_callbacks: Dict[ExtensionCallbackType, List[CallbackFunctionWithMetadata]] = {}
 
-    def add_extension(self, ext: ExtensionBase):
-        self.extensions.append(ext)
+    def add_extension(self, extension: ExtensionBase):
+        self._extensions.append(extension)
+        self._regenerate_ordered_callbacks()
 
-        self._callbacks.clear()
+    def _regenerate_ordered_callbacks(self):
+        """Regenerates self._ordered_callbacks. Intended to be called each time a new extension is added."""
+        self._ordered_callbacks = {}
 
-        for ext in self.extensions:
-            for inj_info in ext.injections:
-                if inj_info.type == "callback":
-                    if inj_info.name not in self._callbacks:
-                        self._callbacks[inj_info.name] = CallbackInjectionPoint()
-                    self._callbacks[inj_info.name].add(inj_info.function, inj_info.order)
+        # Fill the ordered callbacks dictionary.
+        for extension in self._extensions:
+            for callback_type, callbacks in extension.get_callbacks().items():
+                if callback_type not in self._ordered_callbacks:
+                    self._ordered_callbacks[callback_type] = []
+                self._ordered_callbacks[callback_type].extend(callbacks)
 
-                else:
-                    raise Exception(f"Unsupported injection type: {inj_info.type}")
+        # Sort each callback list.
+        for callback_type, callbacks in self._ordered_callbacks.items():
+            self._ordered_callbacks[callback_type] = sorted(callbacks, key=lambda x: x.metadata.order)
 
-    def call_callback(self, name: str, *args, **kwargs):
+    def run_callback(self, callback_type: ExtensionCallbackType, ctx: DenoiseContext):
         # TODO: add to patchers too?
         # and if so, should it be only in beginning of function or in for loop
         if self._is_canceled and self._is_canceled():
             raise CanceledException
 
-        if name in self._callbacks:
-            self._callbacks[name](*args, **kwargs)
+        callbacks = self._ordered_callbacks.get(callback_type, [])
+        for cb in callbacks:
+            cb.function(ctx)
 
     @contextmanager
     def patch_extensions(self, context: DenoiseContext):
         with ExitStack() as exit_stack:
-            for ext in self.extensions:
+            for ext in self._extensions:
                 exit_stack.enter_context(ext.patch_extension(context))
 
             yield None