Remove unused layer_key property from LoRALayerBase.

Need to figure out why the rect isn't reset when the entity is reset. Probably just needs some special handling.

invokeai/app/api/routers/session_queue.py

@@ -11,6 +11,7 @@ from invokeai.app.services.session_queue.session_queue_common import (
     Batch,
     BatchStatus,
     CancelByBatchIDsResult,
+    CancelByDestinationResult,
     ClearResult,
     EnqueueBatchResult,
     PruneResult,
@@ -105,6 +106,21 @@ async def cancel_by_batch_ids(
     return ApiDependencies.invoker.services.session_queue.cancel_by_batch_ids(queue_id=queue_id, batch_ids=batch_ids)
 
 
+@session_queue_router.put(
+    "/{queue_id}/cancel_by_destination",
+    operation_id="cancel_by_destination",
+    responses={200: {"model": CancelByBatchIDsResult}},
+)
+async def cancel_by_destination(
+    queue_id: str = Path(description="The queue id to perform this operation on"),
+    destination: str = Query(description="The destination to cancel all queue items for"),
+) -> CancelByDestinationResult:
+    """Immediately cancels all queue items with the given origin"""
+    return ApiDependencies.invoker.services.session_queue.cancel_by_destination(
+        queue_id=queue_id, destination=destination
+    )
+
+
 @session_queue_router.put(
     "/{queue_id}/clear",
     operation_id="clear",

invokeai/app/invocations/baseinvocation.py

@@ -20,7 +20,6 @@ from typing import (
     Type,
     TypeVar,
     Union,
-    cast,
 )
 
 import semver
@@ -80,7 +79,7 @@ class UIConfigBase(BaseModel):
     version: str = Field(
         description='The node\'s version. Should be a valid semver string e.g. "1.0.0" or "3.8.13".',
     )
-    node_pack: Optional[str] = Field(default=None, description="Whether or not this is a custom node")
+    node_pack: str = Field(description="The node pack that this node belongs to, will be 'invokeai' for built-in nodes")
     classification: Classification = Field(default=Classification.Stable, description="The node's classification")
 
     model_config = ConfigDict(
@@ -230,18 +229,16 @@ class BaseInvocation(ABC, BaseModel):
     @staticmethod
     def json_schema_extra(schema: dict[str, Any], model_class: Type[BaseInvocation]) -> None:
         """Adds various UI-facing attributes to the invocation's OpenAPI schema."""
-        uiconfig = cast(UIConfigBase | None, getattr(model_class, "UIConfig", None))
-        if uiconfig is not None:
-            if uiconfig.title is not None:
-                schema["title"] = uiconfig.title
-            if uiconfig.tags is not None:
-                schema["tags"] = uiconfig.tags
-            if uiconfig.category is not None:
-                schema["category"] = uiconfig.category
-            if uiconfig.node_pack is not None:
-                schema["node_pack"] = uiconfig.node_pack
-            schema["classification"] = uiconfig.classification
-            schema["version"] = uiconfig.version
+        if title := model_class.UIConfig.title:
+            schema["title"] = title
+        if tags := model_class.UIConfig.tags:
+            schema["tags"] = tags
+        if category := model_class.UIConfig.category:
+            schema["category"] = category
+        if node_pack := model_class.UIConfig.node_pack:
+            schema["node_pack"] = node_pack
+        schema["classification"] = model_class.UIConfig.classification
+        schema["version"] = model_class.UIConfig.version
         if "required" not in schema or not isinstance(schema["required"], list):
             schema["required"] = []
         schema["class"] = "invocation"
@@ -312,7 +309,7 @@ class BaseInvocation(ABC, BaseModel):
         json_schema_extra={"field_kind": FieldKind.NodeAttribute},
     )
 
-    UIConfig: ClassVar[Type[UIConfigBase]]
+    UIConfig: ClassVar[UIConfigBase]
 
     model_config = ConfigDict(
         protected_namespaces=(),
@@ -441,30 +438,25 @@ def invocation(
         validate_fields(cls.model_fields, invocation_type)
 
         # Add OpenAPI schema extras
-        uiconfig_name = cls.__qualname__ + ".UIConfig"
-        if not hasattr(cls, "UIConfig") or cls.UIConfig.__qualname__ != uiconfig_name:
-            cls.UIConfig = type(uiconfig_name, (UIConfigBase,), {})
-        cls.UIConfig.title = title
-        cls.UIConfig.tags = tags
-        cls.UIConfig.category = category
-        cls.UIConfig.classification = classification
-
-        # Grab the node pack's name from the module name, if it's a custom node
-        is_custom_node = cls.__module__.rsplit(".", 1)[0] == "invokeai.app.invocations"
-        if is_custom_node:
-            cls.UIConfig.node_pack = cls.__module__.split(".")[0]
-        else:
-            cls.UIConfig.node_pack = None
+        uiconfig: dict[str, Any] = {}
+        uiconfig["title"] = title
+        uiconfig["tags"] = tags
+        uiconfig["category"] = category
+        uiconfig["classification"] = classification
+        # The node pack is the module name - will be "invokeai" for built-in nodes
+        uiconfig["node_pack"] = cls.__module__.split(".")[0]
 
         if version is not None:
             try:
                 semver.Version.parse(version)
             except ValueError as e:
                 raise InvalidVersionError(f'Invalid version string for node "{invocation_type}": "{version}"') from e
-            cls.UIConfig.version = version
+            uiconfig["version"] = version
         else:
             logger.warn(f'No version specified for node "{invocation_type}", using "1.0.0"')
-            cls.UIConfig.version = "1.0.0"
+            uiconfig["version"] = "1.0.0"
+
+        cls.UIConfig = UIConfigBase(**uiconfig)
 
         if use_cache is not None:
             cls.model_fields["use_cache"].default = use_cache

invokeai/app/invocations/compel.py

@@ -19,7 +19,8 @@ from invokeai.app.invocations.model import CLIPField
 from invokeai.app.invocations.primitives import ConditioningOutput
 from invokeai.app.services.shared.invocation_context import InvocationContext
 from invokeai.app.util.ti_utils import generate_ti_list
-from invokeai.backend.lora import LoRAModelRaw
+from invokeai.backend.lora.lora_model_raw import LoRAModelRaw
+from invokeai.backend.lora.lora_patcher import LoraPatcher
 from invokeai.backend.model_patcher import ModelPatcher
 from invokeai.backend.stable_diffusion.diffusion.conditioning_data import (
     BasicConditioningInfo,
@@ -82,9 +83,10 @@ class CompelInvocation(BaseInvocation):
             # apply all patches while the model is on the target device
             text_encoder_info.model_on_device() as (cached_weights, text_encoder),
             tokenizer_info as tokenizer,
-            ModelPatcher.apply_lora_text_encoder(
-                text_encoder,
-                loras=_lora_loader(),
+            LoraPatcher.apply_lora_patches(
+                model=text_encoder,
+                patches=_lora_loader(),
+                prefix="lora_te_",
                 cached_weights=cached_weights,
             ),
             # Apply CLIP Skip after LoRA to prevent LoRA application from failing on skipped layers.
@@ -177,9 +179,9 @@ class SDXLPromptInvocationBase:
             # apply all patches while the model is on the target device
             text_encoder_info.model_on_device() as (cached_weights, text_encoder),
             tokenizer_info as tokenizer,
-            ModelPatcher.apply_lora(
+            LoraPatcher.apply_lora_patches(
                 text_encoder,
-                loras=_lora_loader(),
+                patches=_lora_loader(),
                 prefix=lora_prefix,
                 cached_weights=cached_weights,
             ),

invokeai/app/invocations/denoise_latents.py

@@ -36,7 +36,8 @@ from invokeai.app.invocations.t2i_adapter import T2IAdapterField
 from invokeai.app.services.shared.invocation_context import InvocationContext
 from invokeai.app.util.controlnet_utils import prepare_control_image
 from invokeai.backend.ip_adapter.ip_adapter import IPAdapter
-from invokeai.backend.lora import LoRAModelRaw
+from invokeai.backend.lora.lora_model_raw import LoRAModelRaw
+from invokeai.backend.lora.lora_patcher import LoraPatcher
 from invokeai.backend.model_manager import BaseModelType, ModelVariantType
 from invokeai.backend.model_patcher import ModelPatcher
 from invokeai.backend.stable_diffusion import PipelineIntermediateState
@@ -979,9 +980,10 @@ class DenoiseLatentsInvocation(BaseInvocation):
             ModelPatcher.apply_freeu(unet, self.unet.freeu_config),
             SeamlessExt.static_patch_model(unet, self.unet.seamless_axes),  # FIXME
             # Apply the LoRA after unet has been moved to its target device for faster patching.
-            ModelPatcher.apply_lora_unet(
-                unet,
-                loras=_lora_loader(),
+            LoraPatcher.apply_lora_patches(
+                model=unet,
+                patches=_lora_loader(),
+                prefix="lora_unet_",
                 cached_weights=cached_weights,
             ),
         ):

invokeai/app/invocations/flux_denoise.py

@@ -1,4 +1,4 @@
-from typing import Callable, Optional
+from typing import Callable, Iterator, Optional, Tuple
 
 import torch
 import torchvision.transforms as tv_transforms
@@ -29,6 +29,8 @@ from invokeai.backend.flux.sampling_utils import (
     pack,
     unpack,
 )
+from invokeai.backend.lora.lora_model_raw import LoRAModelRaw
+from invokeai.backend.lora.lora_patcher import LoraPatcher
 from invokeai.backend.stable_diffusion.diffusers_pipeline import PipelineIntermediateState
 from invokeai.backend.stable_diffusion.diffusion.conditioning_data import FLUXConditioningInfo
 from invokeai.backend.util.devices import TorchDevice
@@ -187,7 +189,16 @@ class FluxDenoiseInvocation(BaseInvocation, WithMetadata, WithBoard):
                 noise=noise,
             )
 
-        with transformer_info as transformer:
+        with (
+            transformer_info.model_on_device() as (cached_weights, transformer),
+            # Apply the LoRA after transformer has been moved to its target device for faster patching.
+            LoraPatcher.apply_lora_patches(
+                model=transformer,
+                patches=self._lora_iterator(context),
+                prefix="",
+                cached_weights=cached_weights,
+            ),
+        ):
             assert isinstance(transformer, Flux)
 
             x = denoise(
@@ -241,6 +252,13 @@ class FluxDenoiseInvocation(BaseInvocation, WithMetadata, WithBoard):
         # `latents`.
         return mask.expand_as(latents)
 
+    def _lora_iterator(self, context: InvocationContext) -> Iterator[Tuple[LoRAModelRaw, float]]:
+        for lora in self.transformer.loras:
+            lora_info = context.models.load(lora.lora)
+            assert isinstance(lora_info.model, LoRAModelRaw)
+            yield (lora_info.model, lora.weight)
+            del lora_info
+
     def _build_step_callback(self, context: InvocationContext) -> Callable[[PipelineIntermediateState], None]:
         def step_callback(state: PipelineIntermediateState) -> None:
             state.latents = unpack(state.latents.float(), self.height, self.width).squeeze()

invokeai/app/invocations/flux_lora_loader.py

@@ -0,0 +1,53 @@
+from invokeai.app.invocations.baseinvocation import BaseInvocation, BaseInvocationOutput, invocation, invocation_output
+from invokeai.app.invocations.fields import FieldDescriptions, Input, InputField, OutputField, UIType
+from invokeai.app.invocations.model import LoRAField, ModelIdentifierField, TransformerField
+from invokeai.app.services.shared.invocation_context import InvocationContext
+
+
+@invocation_output("flux_lora_loader_output")
+class FluxLoRALoaderOutput(BaseInvocationOutput):
+    """FLUX LoRA Loader Output"""
+
+    transformer: TransformerField = OutputField(
+        default=None, description=FieldDescriptions.transformer, title="FLUX Transformer"
+    )
+
+
+@invocation(
+    "flux_lora_loader",
+    title="FLUX LoRA",
+    tags=["lora", "model", "flux"],
+    category="model",
+    version="1.0.0",
+)
+class FluxLoRALoaderInvocation(BaseInvocation):
+    """Apply a LoRA model to a FLUX transformer."""
+
+    lora: ModelIdentifierField = InputField(
+        description=FieldDescriptions.lora_model, title="LoRA", ui_type=UIType.LoRAModel
+    )
+    weight: float = InputField(default=0.75, description=FieldDescriptions.lora_weight)
+    transformer: TransformerField = InputField(
+        description=FieldDescriptions.transformer,
+        input=Input.Connection,
+        title="FLUX Transformer",
+    )
+
+    def invoke(self, context: InvocationContext) -> FluxLoRALoaderOutput:
+        lora_key = self.lora.key
+
+        if not context.models.exists(lora_key):
+            raise ValueError(f"Unknown lora: {lora_key}!")
+
+        if any(lora.lora.key == lora_key for lora in self.transformer.loras):
+            raise Exception(f'LoRA "{lora_key}" already applied to transformer.')
+
+        transformer = self.transformer.model_copy(deep=True)
+        transformer.loras.append(
+            LoRAField(
+                lora=self.lora,
+                weight=self.weight,
+            )
+        )
+
+        return FluxLoRALoaderOutput(transformer=transformer)

invokeai/app/invocations/image.py

@@ -6,13 +6,19 @@ import cv2
 import numpy
 from PIL import Image, ImageChops, ImageFilter, ImageOps
 
-from invokeai.app.invocations.baseinvocation import BaseInvocation, Classification, invocation
+from invokeai.app.invocations.baseinvocation import (
+    BaseInvocation,
+    Classification,
+    invocation,
+    invocation_output,
+)
 from invokeai.app.invocations.constants import IMAGE_MODES
 from invokeai.app.invocations.fields import (
     ColorField,
     FieldDescriptions,
     ImageField,
     InputField,
+    OutputField,
     WithBoard,
     WithMetadata,
 )
@@ -1007,3 +1013,62 @@ class MaskFromIDInvocation(BaseInvocation, WithMetadata, WithBoard):
         image_dto = context.images.save(image=mask, image_category=ImageCategory.MASK)
 
         return ImageOutput.build(image_dto)
+
+
+@invocation_output("canvas_v2_mask_and_crop_output")
+class CanvasV2MaskAndCropOutput(ImageOutput):
+    offset_x: int = OutputField(description="The x offset of the image, after cropping")
+    offset_y: int = OutputField(description="The y offset of the image, after cropping")
+
+
+@invocation(
+    "canvas_v2_mask_and_crop",
+    title="Canvas V2 Mask and Crop",
+    tags=["image", "mask", "id"],
+    category="image",
+    version="1.0.0",
+    classification=Classification.Prototype,
+)
+class CanvasV2MaskAndCropInvocation(BaseInvocation, WithMetadata, WithBoard):
+    """Handles Canvas V2 image output masking and cropping"""
+
+    source_image: ImageField | None = InputField(
+        default=None,
+        description="The source image onto which the masked generated image is pasted. If omitted, the masked generated image is returned with transparency.",
+    )
+    generated_image: ImageField = InputField(description="The image to apply the mask to")
+    mask: ImageField = InputField(description="The mask to apply")
+    mask_blur: int = InputField(default=0, ge=0, description="The amount to blur the mask by")
+
+    def _prepare_mask(self, mask: Image.Image) -> Image.Image:
+        mask_array = numpy.array(mask)
+        kernel = numpy.ones((self.mask_blur, self.mask_blur), numpy.uint8)
+        dilated_mask_array = cv2.erode(mask_array, kernel, iterations=3)
+        dilated_mask = Image.fromarray(dilated_mask_array)
+        if self.mask_blur > 0:
+            mask = dilated_mask.filter(ImageFilter.GaussianBlur(self.mask_blur))
+        return ImageOps.invert(mask.convert("L"))
+
+    def invoke(self, context: InvocationContext) -> CanvasV2MaskAndCropOutput:
+        mask = self._prepare_mask(context.images.get_pil(self.mask.image_name))
+
+        if self.source_image:
+            generated_image = context.images.get_pil(self.generated_image.image_name)
+            source_image = context.images.get_pil(self.source_image.image_name)
+            source_image.paste(generated_image, (0, 0), mask)
+            image_dto = context.images.save(image=source_image)
+        else:
+            generated_image = context.images.get_pil(self.generated_image.image_name)
+            generated_image.putalpha(mask)
+            image_dto = context.images.save(image=generated_image)
+
+        # bbox = image.getbbox()
+        # image = image.crop(bbox)
+
+        return CanvasV2MaskAndCropOutput(
+            image=ImageField(image_name=image_dto.image_name),
+            offset_x=0,
+            offset_y=0,
+            width=image_dto.width,
+            height=image_dto.height,
+        )

invokeai/app/invocations/model.py

@@ -69,6 +69,7 @@ class CLIPField(BaseModel):
 
 class TransformerField(BaseModel):
     transformer: ModelIdentifierField = Field(description="Info to load Transformer submodel")
+    loras: List[LoRAField] = Field(description="LoRAs to apply on model loading")
 
 
 class T5EncoderField(BaseModel):
@@ -202,7 +203,7 @@ class FluxModelLoaderInvocation(BaseInvocation):
         assert isinstance(transformer_config, CheckpointConfigBase)
 
         return FluxModelLoaderOutput(
-            transformer=TransformerField(transformer=transformer),
+            transformer=TransformerField(transformer=transformer, loras=[]),
             clip=CLIPField(tokenizer=tokenizer, text_encoder=clip_encoder, loras=[], skipped_layers=0),
             t5_encoder=T5EncoderField(tokenizer=tokenizer2, text_encoder=t5_encoder),
             vae=VAEField(vae=vae),

invokeai/app/invocations/tiled_multi_diffusion_denoise_latents.py

@@ -22,8 +22,8 @@ from invokeai.app.invocations.fields import (
 from invokeai.app.invocations.model import UNetField
 from invokeai.app.invocations.primitives import LatentsOutput
 from invokeai.app.services.shared.invocation_context import InvocationContext
-from invokeai.backend.lora import LoRAModelRaw
-from invokeai.backend.model_patcher import ModelPatcher
+from invokeai.backend.lora.lora_model_raw import LoRAModelRaw
+from invokeai.backend.lora.lora_patcher import LoraPatcher
 from invokeai.backend.stable_diffusion.diffusers_pipeline import ControlNetData, PipelineIntermediateState
 from invokeai.backend.stable_diffusion.multi_diffusion_pipeline import (
     MultiDiffusionPipeline,
@@ -204,7 +204,11 @@ class TiledMultiDiffusionDenoiseLatents(BaseInvocation):
         # Load the UNet model.
         unet_info = context.models.load(self.unet.unet)
 
-        with ExitStack() as exit_stack, unet_info as unet, ModelPatcher.apply_lora_unet(unet, _lora_loader()):
+        with (
+            ExitStack() as exit_stack,
+            unet_info as unet,
+            LoraPatcher.apply_lora_patches(model=unet, patches=_lora_loader(), prefix="lora_unet_"),
+        ):
             assert isinstance(unet, UNet2DConditionModel)
             latents = latents.to(device=unet.device, dtype=unet.dtype)
             if noise is not None:

invokeai/app/services/events/events_common.py

@@ -88,6 +88,8 @@ class QueueItemEventBase(QueueEventBase):
 
     item_id: int = Field(description="The ID of the queue item")
     batch_id: str = Field(description="The ID of the queue batch")
+    origin: str | None = Field(default=None, description="The origin of the queue item")
+    destination: str | None = Field(default=None, description="The destination of the queue item")
 
 
 class InvocationEventBase(QueueItemEventBase):
@@ -95,8 +97,6 @@ class InvocationEventBase(QueueItemEventBase):
 
     session_id: str = Field(description="The ID of the session (aka graph execution state)")
     queue_id: str = Field(description="The ID of the queue")
-    item_id: int = Field(description="The ID of the queue item")
-    batch_id: str = Field(description="The ID of the queue batch")
     session_id: str = Field(description="The ID of the session (aka graph execution state)")
     invocation: AnyInvocation = Field(description="The ID of the invocation")
     invocation_source_id: str = Field(description="The ID of the prepared invocation's source node")
@@ -114,6 +114,8 @@ class InvocationStartedEvent(InvocationEventBase):
             queue_id=queue_item.queue_id,
             item_id=queue_item.item_id,
             batch_id=queue_item.batch_id,
+            origin=queue_item.origin,
+            destination=queue_item.destination,
             session_id=queue_item.session_id,
             invocation=invocation,
             invocation_source_id=queue_item.session.prepared_source_mapping[invocation.id],
@@ -147,6 +149,8 @@ class InvocationDenoiseProgressEvent(InvocationEventBase):
             queue_id=queue_item.queue_id,
             item_id=queue_item.item_id,
             batch_id=queue_item.batch_id,
+            origin=queue_item.origin,
+            destination=queue_item.destination,
             session_id=queue_item.session_id,
             invocation=invocation,
             invocation_source_id=queue_item.session.prepared_source_mapping[invocation.id],
@@ -184,6 +188,8 @@ class InvocationCompleteEvent(InvocationEventBase):
             queue_id=queue_item.queue_id,
             item_id=queue_item.item_id,
             batch_id=queue_item.batch_id,
+            origin=queue_item.origin,
+            destination=queue_item.destination,
             session_id=queue_item.session_id,
             invocation=invocation,
             invocation_source_id=queue_item.session.prepared_source_mapping[invocation.id],
@@ -216,6 +222,8 @@ class InvocationErrorEvent(InvocationEventBase):
             queue_id=queue_item.queue_id,
             item_id=queue_item.item_id,
             batch_id=queue_item.batch_id,
+            origin=queue_item.origin,
+            destination=queue_item.destination,
             session_id=queue_item.session_id,
             invocation=invocation,
             invocation_source_id=queue_item.session.prepared_source_mapping[invocation.id],
@@ -253,6 +261,8 @@ class QueueItemStatusChangedEvent(QueueItemEventBase):
             queue_id=queue_item.queue_id,
             item_id=queue_item.item_id,
             batch_id=queue_item.batch_id,
+            origin=queue_item.origin,
+            destination=queue_item.destination,
             session_id=queue_item.session_id,
             status=queue_item.status,
             error_type=queue_item.error_type,
@@ -279,12 +289,14 @@ class BatchEnqueuedEvent(QueueEventBase):
         description="The number of invocations initially requested to be enqueued (may be less than enqueued if queue was full)"
     )
     priority: int = Field(description="The priority of the batch")
+    origin: str | None = Field(default=None, description="The origin of the batch")
 
     @classmethod
     def build(cls, enqueue_result: EnqueueBatchResult) -> "BatchEnqueuedEvent":
         return cls(
             queue_id=enqueue_result.queue_id,
             batch_id=enqueue_result.batch.batch_id,
+            origin=enqueue_result.batch.origin,
             enqueued=enqueue_result.enqueued,
             requested=enqueue_result.requested,
             priority=enqueue_result.priority,

invokeai/app/services/session_queue/session_queue_base.py

@@ -6,6 +6,7 @@ from invokeai.app.services.session_queue.session_queue_common import (
     Batch,
     BatchStatus,
     CancelByBatchIDsResult,
+    CancelByDestinationResult,
     CancelByQueueIDResult,
     ClearResult,
     EnqueueBatchResult,
@@ -95,6 +96,11 @@ class SessionQueueBase(ABC):
         """Cancels all queue items with matching batch IDs"""
         pass
 
+    @abstractmethod
+    def cancel_by_destination(self, queue_id: str, destination: str) -> CancelByDestinationResult:
+        """Cancels all queue items with the given batch destination"""
+        pass
+
     @abstractmethod
     def cancel_by_queue_id(self, queue_id: str) -> CancelByQueueIDResult:
         """Cancels all queue items with matching queue ID"""

invokeai/app/services/session_queue/session_queue_common.py

@@ -77,6 +77,14 @@ BatchDataCollection: TypeAlias = list[list[BatchDatum]]
 
 class Batch(BaseModel):
     batch_id: str = Field(default_factory=uuid_string, description="The ID of the batch")
+    origin: str | None = Field(
+        default=None,
+        description="The origin of this queue item. This data is used by the frontend to determine how to handle results.",
+    )
+    destination: str | None = Field(
+        default=None,
+        description="The origin of this queue item. This data is used by the frontend to determine how to handle results",
+    )
     data: Optional[BatchDataCollection] = Field(default=None, description="The batch data collection.")
     graph: Graph = Field(description="The graph to initialize the session with")
     workflow: Optional[WorkflowWithoutID] = Field(
@@ -195,6 +203,14 @@ class SessionQueueItemWithoutGraph(BaseModel):
     status: QUEUE_ITEM_STATUS = Field(default="pending", description="The status of this queue item")
     priority: int = Field(default=0, description="The priority of this queue item")
     batch_id: str = Field(description="The ID of the batch associated with this queue item")
+    origin: str | None = Field(
+        default=None,
+        description="The origin of this queue item. This data is used by the frontend to determine how to handle results.",
+    )
+    destination: str | None = Field(
+        default=None,
+        description="The origin of this queue item. This data is used by the frontend to determine how to handle results",
+    )
     session_id: str = Field(
         description="The ID of the session associated with this queue item. The session doesn't exist in graph_executions until the queue item is executed."
     )
@@ -294,6 +310,8 @@ class SessionQueueStatus(BaseModel):
 class BatchStatus(BaseModel):
     queue_id: str = Field(..., description="The ID of the queue")
     batch_id: str = Field(..., description="The ID of the batch")
+    origin: str | None = Field(..., description="The origin of the batch")
+    destination: str | None = Field(..., description="The destination of the batch")
     pending: int = Field(..., description="Number of queue items with status 'pending'")
     in_progress: int = Field(..., description="Number of queue items with status 'in_progress'")
     completed: int = Field(..., description="Number of queue items with status 'complete'")
@@ -328,6 +346,12 @@ class CancelByBatchIDsResult(BaseModel):
     canceled: int = Field(..., description="Number of queue items canceled")
 
 
+class CancelByDestinationResult(CancelByBatchIDsResult):
+    """Result of canceling by a destination"""
+
+    pass
+
+
 class CancelByQueueIDResult(CancelByBatchIDsResult):
     """Result of canceling by queue id"""
 
@@ -433,6 +457,8 @@ class SessionQueueValueToInsert(NamedTuple):
     field_values: Optional[str]  # field_values json
     priority: int  # priority
     workflow: Optional[str]  # workflow json
+    origin: str | None
+    destination: str | None
 
 
 ValuesToInsert: TypeAlias = list[SessionQueueValueToInsert]
@@ -453,6 +479,8 @@ def prepare_values_to_insert(queue_id: str, batch: Batch, priority: int, max_new
                 json.dumps(field_values, default=to_jsonable_python) if field_values else None,  # field_values (json)
                 priority,  # priority
                 json.dumps(workflow, default=to_jsonable_python) if workflow else None,  # workflow (json)
+                batch.origin,  # origin
+                batch.destination,  # destination
             )
         )
     return values_to_insert

invokeai/app/services/session_queue/session_queue_sqlite.py

@@ -10,6 +10,7 @@ from invokeai.app.services.session_queue.session_queue_common import (
     Batch,
     BatchStatus,
     CancelByBatchIDsResult,
+    CancelByDestinationResult,
     CancelByQueueIDResult,
     ClearResult,
     EnqueueBatchResult,
@@ -127,8 +128,8 @@ class SqliteSessionQueue(SessionQueueBase):
 
             self.__cursor.executemany(
                 """--sql
-                INSERT INTO session_queue (queue_id, session, session_id, batch_id, field_values, priority, workflow)
-                VALUES (?, ?, ?, ?, ?, ?, ?)
+                INSERT INTO session_queue (queue_id, session, session_id, batch_id, field_values, priority, workflow, origin, destination)
+                VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)
                 """,
                 values_to_insert,
             )
@@ -417,11 +418,7 @@ class SqliteSessionQueue(SessionQueueBase):
             )
             self.__conn.commit()
             if current_queue_item is not None and current_queue_item.batch_id in batch_ids:
-                batch_status = self.get_batch_status(queue_id=queue_id, batch_id=current_queue_item.batch_id)
-                queue_status = self.get_queue_status(queue_id=queue_id)
-                self.__invoker.services.events.emit_queue_item_status_changed(
-                    current_queue_item, batch_status, queue_status
-                )
+                self._set_queue_item_status(current_queue_item.item_id, "canceled")
         except Exception:
             self.__conn.rollback()
             raise
@@ -429,6 +426,46 @@ class SqliteSessionQueue(SessionQueueBase):
             self.__lock.release()
         return CancelByBatchIDsResult(canceled=count)
 
+    def cancel_by_destination(self, queue_id: str, destination: str) -> CancelByDestinationResult:
+        try:
+            current_queue_item = self.get_current(queue_id)
+            self.__lock.acquire()
+            where = """--sql
+                WHERE
+                  queue_id == ?
+                  AND destination == ?
+                  AND status != 'canceled'
+                  AND status != 'completed'
+                  AND status != 'failed'
+                """
+            params = (queue_id, destination)
+            self.__cursor.execute(
+                f"""--sql
+                SELECT COUNT(*)
+                FROM session_queue
+                {where};
+                """,
+                params,
+            )
+            count = self.__cursor.fetchone()[0]
+            self.__cursor.execute(
+                f"""--sql
+                UPDATE session_queue
+                SET status = 'canceled'
+                {where};
+                """,
+                params,
+            )
+            self.__conn.commit()
+            if current_queue_item is not None and current_queue_item.destination == destination:
+                self._set_queue_item_status(current_queue_item.item_id, "canceled")
+        except Exception:
+            self.__conn.rollback()
+            raise
+        finally:
+            self.__lock.release()
+        return CancelByDestinationResult(canceled=count)
+
     def cancel_by_queue_id(self, queue_id: str) -> CancelByQueueIDResult:
         try:
             current_queue_item = self.get_current(queue_id)
@@ -541,7 +578,9 @@ class SqliteSessionQueue(SessionQueueBase):
                     started_at,
                     session_id,
                     batch_id,
-                    queue_id
+                    queue_id,
+                    origin,
+                    destination
                 FROM session_queue
                 WHERE queue_id = ?
             """
@@ -621,7 +660,7 @@ class SqliteSessionQueue(SessionQueueBase):
             self.__lock.acquire()
             self.__cursor.execute(
                 """--sql
-                SELECT status, count(*)
+                SELECT status, count(*), origin, destination
                 FROM session_queue
                 WHERE
                   queue_id = ?
@@ -633,6 +672,8 @@ class SqliteSessionQueue(SessionQueueBase):
             result = cast(list[sqlite3.Row], self.__cursor.fetchall())
             total = sum(row[1] for row in result)
             counts: dict[str, int] = {row[0]: row[1] for row in result}
+            origin = result[0]["origin"] if result else None
+            destination = result[0]["destination"] if result else None
         except Exception:
             self.__conn.rollback()
             raise
@@ -641,6 +682,8 @@ class SqliteSessionQueue(SessionQueueBase):
 
         return BatchStatus(
             batch_id=batch_id,
+            origin=origin,
+            destination=destination,
             queue_id=queue_id,
             pending=counts.get("pending", 0),
             in_progress=counts.get("in_progress", 0),

invokeai/app/services/shared/sqlite/sqlite_util.py

@@ -17,6 +17,7 @@ from invokeai.app.services.shared.sqlite_migrator.migrations.migration_11 import
 from invokeai.app.services.shared.sqlite_migrator.migrations.migration_12 import build_migration_12
 from invokeai.app.services.shared.sqlite_migrator.migrations.migration_13 import build_migration_13
 from invokeai.app.services.shared.sqlite_migrator.migrations.migration_14 import build_migration_14
+from invokeai.app.services.shared.sqlite_migrator.migrations.migration_15 import build_migration_15
 from invokeai.app.services.shared.sqlite_migrator.sqlite_migrator_impl import SqliteMigrator
 
 
@@ -51,6 +52,7 @@ def init_db(config: InvokeAIAppConfig, logger: Logger, image_files: ImageFileSto
     migrator.register_migration(build_migration_12(app_config=config))
     migrator.register_migration(build_migration_13())
     migrator.register_migration(build_migration_14())
+    migrator.register_migration(build_migration_15())
     migrator.run_migrations()
 
     return db

invokeai/app/services/shared/sqlite_migrator/migrations/migration_15.py

@@ -0,0 +1,34 @@
+import sqlite3
+
+from invokeai.app.services.shared.sqlite_migrator.sqlite_migrator_common import Migration
+
+
+class Migration15Callback:
+    def __call__(self, cursor: sqlite3.Cursor) -> None:
+        self._add_origin_col(cursor)
+
+    def _add_origin_col(self, cursor: sqlite3.Cursor) -> None:
+        """
+        - Adds `origin` column to the session queue table.
+        - Adds `destination` column to the session queue table.
+        """
+
+        cursor.execute("ALTER TABLE session_queue ADD COLUMN origin TEXT;")
+        cursor.execute("ALTER TABLE session_queue ADD COLUMN destination TEXT;")
+
+
+def build_migration_15() -> Migration:
+    """
+    Build the migration from database version 14 to 15.
+
+    This migration does the following:
+        - Adds `origin` column to the session queue table.
+        - Adds `destination` column to the session queue table.
+    """
+    migration_15 = Migration(
+        from_version=14,
+        to_version=15,
+        callback=Migration15Callback(),
+    )
+
+    return migration_15

invokeai/backend/lora.py

@@ -1,672 +0,0 @@
-# Copyright (c) 2024 The InvokeAI Development team
-"""LoRA model support."""
-
-import bisect
-from pathlib import Path
-from typing import Dict, List, Optional, Set, Tuple, Union
-
-import torch
-from safetensors.torch import load_file
-from typing_extensions import Self
-
-import invokeai.backend.util.logging as logger
-from invokeai.backend.model_manager import BaseModelType
-from invokeai.backend.raw_model import RawModel
-
-
-class LoRALayerBase:
-    # rank: Optional[int]
-    # alpha: Optional[float]
-    # bias: Optional[torch.Tensor]
-    # layer_key: str
-
-    # @property
-    # def scale(self):
-    #    return self.alpha / self.rank if (self.alpha and self.rank) else 1.0
-
-    def __init__(
-        self,
-        layer_key: str,
-        values: Dict[str, torch.Tensor],
-    ):
-        if "alpha" in values:
-            self.alpha = values["alpha"].item()
-        else:
-            self.alpha = None
-
-        if "bias_indices" in values and "bias_values" in values and "bias_size" in values:
-            self.bias: Optional[torch.Tensor] = torch.sparse_coo_tensor(
-                values["bias_indices"],
-                values["bias_values"],
-                tuple(values["bias_size"]),
-            )
-
-        else:
-            self.bias = None
-
-        self.rank = None  # set in layer implementation
-        self.layer_key = layer_key
-
-    def get_weight(self, orig_weight: torch.Tensor) -> torch.Tensor:
-        raise NotImplementedError()
-
-    def get_bias(self, orig_bias: torch.Tensor) -> Optional[torch.Tensor]:
-        return self.bias
-
-    def get_parameters(self, orig_module: torch.nn.Module) -> Dict[str, torch.Tensor]:
-        params = {"weight": self.get_weight(orig_module.weight)}
-        bias = self.get_bias(orig_module.bias)
-        if bias is not None:
-            params["bias"] = bias
-        return params
-
-    def calc_size(self) -> int:
-        model_size = 0
-        for val in [self.bias]:
-            if val is not None:
-                model_size += val.nelement() * val.element_size()
-        return model_size
-
-    def to(self, device: Optional[torch.device] = None, dtype: Optional[torch.dtype] = None) -> None:
-        if self.bias is not None:
-            self.bias = self.bias.to(device=device, dtype=dtype)
-
-    def check_keys(self, values: Dict[str, torch.Tensor], known_keys: Set[str]):
-        """Log a warning if values contains unhandled keys."""
-        # {"alpha", "bias_indices", "bias_values", "bias_size"} are hard-coded, because they are handled by
-        # `LoRALayerBase`. Sub-classes should provide the known_keys that they handled.
-        all_known_keys = known_keys | {"alpha", "bias_indices", "bias_values", "bias_size"}
-        unknown_keys = set(values.keys()) - all_known_keys
-        if unknown_keys:
-            logger.warning(
-                f"Unexpected keys found in LoRA/LyCORIS layer, model might work incorrectly! Keys: {unknown_keys}"
-            )
-
-
-# TODO: find and debug lora/locon with bias
-class LoRALayer(LoRALayerBase):
-    # up: torch.Tensor
-    # mid: Optional[torch.Tensor]
-    # down: torch.Tensor
-
-    def __init__(
-        self,
-        layer_key: str,
-        values: Dict[str, torch.Tensor],
-    ):
-        super().__init__(layer_key, values)
-
-        self.up = values["lora_up.weight"]
-        self.down = values["lora_down.weight"]
-        self.mid = values.get("lora_mid.weight", None)
-
-        self.rank = self.down.shape[0]
-        self.check_keys(
-            values,
-            {
-                "lora_up.weight",
-                "lora_down.weight",
-                "lora_mid.weight",
-            },
-        )
-
-    def get_weight(self, orig_weight: torch.Tensor) -> torch.Tensor:
-        if self.mid is not None:
-            up = self.up.reshape(self.up.shape[0], self.up.shape[1])
-            down = self.down.reshape(self.down.shape[0], self.down.shape[1])
-            weight = torch.einsum("m n w h, i m, n j -> i j w h", self.mid, up, down)
-        else:
-            weight = self.up.reshape(self.up.shape[0], -1) @ self.down.reshape(self.down.shape[0], -1)
-
-        return weight
-
-    def calc_size(self) -> int:
-        model_size = super().calc_size()
-        for val in [self.up, self.mid, self.down]:
-            if val is not None:
-                model_size += val.nelement() * val.element_size()
-        return model_size
-
-    def to(self, device: Optional[torch.device] = None, dtype: Optional[torch.dtype] = None) -> None:
-        super().to(device=device, dtype=dtype)
-
-        self.up = self.up.to(device=device, dtype=dtype)
-        self.down = self.down.to(device=device, dtype=dtype)
-
-        if self.mid is not None:
-            self.mid = self.mid.to(device=device, dtype=dtype)
-
-
-class LoHALayer(LoRALayerBase):
-    # w1_a: torch.Tensor
-    # w1_b: torch.Tensor
-    # w2_a: torch.Tensor
-    # w2_b: torch.Tensor
-    # t1: Optional[torch.Tensor] = None
-    # t2: Optional[torch.Tensor] = None
-
-    def __init__(self, layer_key: str, values: Dict[str, torch.Tensor]):
-        super().__init__(layer_key, values)
-
-        self.w1_a = values["hada_w1_a"]
-        self.w1_b = values["hada_w1_b"]
-        self.w2_a = values["hada_w2_a"]
-        self.w2_b = values["hada_w2_b"]
-        self.t1 = values.get("hada_t1", None)
-        self.t2 = values.get("hada_t2", None)
-
-        self.rank = self.w1_b.shape[0]
-        self.check_keys(
-            values,
-            {
-                "hada_w1_a",
-                "hada_w1_b",
-                "hada_w2_a",
-                "hada_w2_b",
-                "hada_t1",
-                "hada_t2",
-            },
-        )
-
-    def get_weight(self, orig_weight: torch.Tensor) -> torch.Tensor:
-        if self.t1 is None:
-            weight: torch.Tensor = (self.w1_a @ self.w1_b) * (self.w2_a @ self.w2_b)
-
-        else:
-            rebuild1 = torch.einsum("i j k l, j r, i p -> p r k l", self.t1, self.w1_b, self.w1_a)
-            rebuild2 = torch.einsum("i j k l, j r, i p -> p r k l", self.t2, self.w2_b, self.w2_a)
-            weight = rebuild1 * rebuild2
-
-        return weight
-
-    def calc_size(self) -> int:
-        model_size = super().calc_size()
-        for val in [self.w1_a, self.w1_b, self.w2_a, self.w2_b, self.t1, self.t2]:
-            if val is not None:
-                model_size += val.nelement() * val.element_size()
-        return model_size
-
-    def to(self, device: Optional[torch.device] = None, dtype: Optional[torch.dtype] = None) -> None:
-        super().to(device=device, dtype=dtype)
-
-        self.w1_a = self.w1_a.to(device=device, dtype=dtype)
-        self.w1_b = self.w1_b.to(device=device, dtype=dtype)
-        if self.t1 is not None:
-            self.t1 = self.t1.to(device=device, dtype=dtype)
-
-        self.w2_a = self.w2_a.to(device=device, dtype=dtype)
-        self.w2_b = self.w2_b.to(device=device, dtype=dtype)
-        if self.t2 is not None:
-            self.t2 = self.t2.to(device=device, dtype=dtype)
-
-
-class LoKRLayer(LoRALayerBase):
-    # w1: Optional[torch.Tensor] = None
-    # w1_a: Optional[torch.Tensor] = None
-    # w1_b: Optional[torch.Tensor] = None
-    # w2: Optional[torch.Tensor] = None
-    # w2_a: Optional[torch.Tensor] = None
-    # w2_b: Optional[torch.Tensor] = None
-    # t2: Optional[torch.Tensor] = None
-
-    def __init__(
-        self,
-        layer_key: str,
-        values: Dict[str, torch.Tensor],
-    ):
-        super().__init__(layer_key, values)
-
-        self.w1 = values.get("lokr_w1", None)
-        if self.w1 is None:
-            self.w1_a = values["lokr_w1_a"]
-            self.w1_b = values["lokr_w1_b"]
-        else:
-            self.w1_b = None
-            self.w1_a = None
-
-        self.w2 = values.get("lokr_w2", None)
-        if self.w2 is None:
-            self.w2_a = values["lokr_w2_a"]
-            self.w2_b = values["lokr_w2_b"]
-        else:
-            self.w2_a = None
-            self.w2_b = None
-
-        self.t2 = values.get("lokr_t2", None)
-
-        if self.w1_b is not None:
-            self.rank = self.w1_b.shape[0]
-        elif self.w2_b is not None:
-            self.rank = self.w2_b.shape[0]
-        else:
-            self.rank = None  # unscaled
-
-        self.check_keys(
-            values,
-            {
-                "lokr_w1",
-                "lokr_w1_a",
-                "lokr_w1_b",
-                "lokr_w2",
-                "lokr_w2_a",
-                "lokr_w2_b",
-                "lokr_t2",
-            },
-        )
-
-    def get_weight(self, orig_weight: torch.Tensor) -> torch.Tensor:
-        w1: Optional[torch.Tensor] = self.w1
-        if w1 is None:
-            assert self.w1_a is not None
-            assert self.w1_b is not None
-            w1 = self.w1_a @ self.w1_b
-
-        w2 = self.w2
-        if w2 is None:
-            if self.t2 is None:
-                assert self.w2_a is not None
-                assert self.w2_b is not None
-                w2 = self.w2_a @ self.w2_b
-            else:
-                w2 = torch.einsum("i j k l, i p, j r -> p r k l", self.t2, self.w2_a, self.w2_b)
-
-        if len(w2.shape) == 4:
-            w1 = w1.unsqueeze(2).unsqueeze(2)
-        w2 = w2.contiguous()
-        assert w1 is not None
-        assert w2 is not None
-        weight = torch.kron(w1, w2)
-
-        return weight
-
-    def calc_size(self) -> int:
-        model_size = super().calc_size()
-        for val in [self.w1, self.w1_a, self.w1_b, self.w2, self.w2_a, self.w2_b, self.t2]:
-            if val is not None:
-                model_size += val.nelement() * val.element_size()
-        return model_size
-
-    def to(self, device: Optional[torch.device] = None, dtype: Optional[torch.dtype] = None) -> None:
-        super().to(device=device, dtype=dtype)
-
-        if self.w1 is not None:
-            self.w1 = self.w1.to(device=device, dtype=dtype)
-        else:
-            assert self.w1_a is not None
-            assert self.w1_b is not None
-            self.w1_a = self.w1_a.to(device=device, dtype=dtype)
-            self.w1_b = self.w1_b.to(device=device, dtype=dtype)
-
-        if self.w2 is not None:
-            self.w2 = self.w2.to(device=device, dtype=dtype)
-        else:
-            assert self.w2_a is not None
-            assert self.w2_b is not None
-            self.w2_a = self.w2_a.to(device=device, dtype=dtype)
-            self.w2_b = self.w2_b.to(device=device, dtype=dtype)
-
-        if self.t2 is not None:
-            self.t2 = self.t2.to(device=device, dtype=dtype)
-
-
-class FullLayer(LoRALayerBase):
-    # bias handled in LoRALayerBase(calc_size, to)
-    # weight: torch.Tensor
-    # bias: Optional[torch.Tensor]
-
-    def __init__(
-        self,
-        layer_key: str,
-        values: Dict[str, torch.Tensor],
-    ):
-        super().__init__(layer_key, values)
-
-        self.weight = values["diff"]
-        self.bias = values.get("diff_b", None)
-
-        self.rank = None  # unscaled
-        self.check_keys(values, {"diff", "diff_b"})
-
-    def get_weight(self, orig_weight: torch.Tensor) -> torch.Tensor:
-        return self.weight
-
-    def calc_size(self) -> int:
-        model_size = super().calc_size()
-        model_size += self.weight.nelement() * self.weight.element_size()
-        return model_size
-
-    def to(self, device: Optional[torch.device] = None, dtype: Optional[torch.dtype] = None) -> None:
-        super().to(device=device, dtype=dtype)
-
-        self.weight = self.weight.to(device=device, dtype=dtype)
-
-
-class IA3Layer(LoRALayerBase):
-    # weight: torch.Tensor
-    # on_input: torch.Tensor
-
-    def __init__(
-        self,
-        layer_key: str,
-        values: Dict[str, torch.Tensor],
-    ):
-        super().__init__(layer_key, values)
-
-        self.weight = values["weight"]
-        self.on_input = values["on_input"]
-
-        self.rank = None  # unscaled
-        self.check_keys(values, {"weight", "on_input"})
-
-    def get_weight(self, orig_weight: torch.Tensor) -> torch.Tensor:
-        weight = self.weight
-        if not self.on_input:
-            weight = weight.reshape(-1, 1)
-        assert orig_weight is not None
-        return orig_weight * weight
-
-    def calc_size(self) -> int:
-        model_size = super().calc_size()
-        model_size += self.weight.nelement() * self.weight.element_size()
-        model_size += self.on_input.nelement() * self.on_input.element_size()
-        return model_size
-
-    def to(self, device: Optional[torch.device] = None, dtype: Optional[torch.dtype] = None):
-        super().to(device=device, dtype=dtype)
-
-        self.weight = self.weight.to(device=device, dtype=dtype)
-        self.on_input = self.on_input.to(device=device, dtype=dtype)
-
-
-class NormLayer(LoRALayerBase):
-    # bias handled in LoRALayerBase(calc_size, to)
-    # weight: torch.Tensor
-    # bias: Optional[torch.Tensor]
-
-    def __init__(
-        self,
-        layer_key: str,
-        values: Dict[str, torch.Tensor],
-    ):
-        super().__init__(layer_key, values)
-
-        self.weight = values["w_norm"]
-        self.bias = values.get("b_norm", None)
-
-        self.rank = None  # unscaled
-        self.check_keys(values, {"w_norm", "b_norm"})
-
-    def get_weight(self, orig_weight: torch.Tensor) -> torch.Tensor:
-        return self.weight
-
-    def calc_size(self) -> int:
-        model_size = super().calc_size()
-        model_size += self.weight.nelement() * self.weight.element_size()
-        return model_size
-
-    def to(self, device: Optional[torch.device] = None, dtype: Optional[torch.dtype] = None) -> None:
-        super().to(device=device, dtype=dtype)
-
-        self.weight = self.weight.to(device=device, dtype=dtype)
-
-
-AnyLoRALayer = Union[LoRALayer, LoHALayer, LoKRLayer, FullLayer, IA3Layer, NormLayer]
-
-
-class LoRAModelRaw(RawModel):  # (torch.nn.Module):
-    _name: str
-    layers: Dict[str, AnyLoRALayer]
-
-    def __init__(
-        self,
-        name: str,
-        layers: Dict[str, AnyLoRALayer],
-    ):
-        self._name = name
-        self.layers = layers
-
-    @property
-    def name(self) -> str:
-        return self._name
-
-    def to(self, device: Optional[torch.device] = None, dtype: Optional[torch.dtype] = None) -> None:
-        # TODO: try revert if exception?
-        for _key, layer in self.layers.items():
-            layer.to(device=device, dtype=dtype)
-
-    def calc_size(self) -> int:
-        model_size = 0
-        for _, layer in self.layers.items():
-            model_size += layer.calc_size()
-        return model_size
-
-    @classmethod
-    def _convert_sdxl_keys_to_diffusers_format(cls, state_dict: Dict[str, torch.Tensor]) -> Dict[str, torch.Tensor]:
-        """Convert the keys of an SDXL LoRA state_dict to diffusers format.
-
-        The input state_dict can be in either Stability AI format or diffusers format. If the state_dict is already in
-        diffusers format, then this function will have no effect.
-
-        This function is adapted from:
-        https://github.com/bmaltais/kohya_ss/blob/2accb1305979ba62f5077a23aabac23b4c37e935/networks/lora_diffusers.py#L385-L409
-
-        Args:
-            state_dict (Dict[str, Tensor]): The SDXL LoRA state_dict.
-
-        Raises:
-            ValueError: If state_dict contains an unrecognized key, or not all keys could be converted.
-
-        Returns:
-            Dict[str, Tensor]: The diffusers-format state_dict.
-        """
-        converted_count = 0  # The number of Stability AI keys converted to diffusers format.
-        not_converted_count = 0  # The number of keys that were not converted.
-
-        # Get a sorted list of Stability AI UNet keys so that we can efficiently search for keys with matching prefixes.
-        # For example, we want to efficiently find `input_blocks_4_1` in the list when searching for
-        # `input_blocks_4_1_proj_in`.
-        stability_unet_keys = list(SDXL_UNET_STABILITY_TO_DIFFUSERS_MAP)
-        stability_unet_keys.sort()
-
-        new_state_dict = {}
-        for full_key, value in state_dict.items():
-            if full_key.startswith("lora_unet_"):
-                search_key = full_key.replace("lora_unet_", "")
-                # Use bisect to find the key in stability_unet_keys that *may* match the search_key's prefix.
-                position = bisect.bisect_right(stability_unet_keys, search_key)
-                map_key = stability_unet_keys[position - 1]
-                # Now, check if the map_key *actually* matches the search_key.
-                if search_key.startswith(map_key):
-                    new_key = full_key.replace(map_key, SDXL_UNET_STABILITY_TO_DIFFUSERS_MAP[map_key])
-                    new_state_dict[new_key] = value
-                    converted_count += 1
-                else:
-                    new_state_dict[full_key] = value
-                    not_converted_count += 1
-            elif full_key.startswith("lora_te1_") or full_key.startswith("lora_te2_"):
-                # The CLIP text encoders have the same keys in both Stability AI and diffusers formats.
-                new_state_dict[full_key] = value
-                continue
-            else:
-                raise ValueError(f"Unrecognized SDXL LoRA key prefix: '{full_key}'.")
-
-        if converted_count > 0 and not_converted_count > 0:
-            raise ValueError(
-                f"The SDXL LoRA could only be partially converted to diffusers format. converted={converted_count},"
-                f" not_converted={not_converted_count}"
-            )
-
-        return new_state_dict
-
-    @classmethod
-    def from_checkpoint(
-        cls,
-        file_path: Union[str, Path],
-        device: Optional[torch.device] = None,
-        dtype: Optional[torch.dtype] = None,
-        base_model: Optional[BaseModelType] = None,
-    ) -> Self:
-        device = device or torch.device("cpu")
-        dtype = dtype or torch.float32
-
-        if isinstance(file_path, str):
-            file_path = Path(file_path)
-
-        model = cls(
-            name=file_path.stem,
-            layers={},
-        )
-
-        if file_path.suffix == ".safetensors":
-            sd = load_file(file_path.absolute().as_posix(), device="cpu")
-        else:
-            sd = torch.load(file_path, map_location="cpu")
-
-        state_dict = cls._group_state(sd)
-
-        if base_model == BaseModelType.StableDiffusionXL:
-            state_dict = cls._convert_sdxl_keys_to_diffusers_format(state_dict)
-
-        for layer_key, values in state_dict.items():
-            # Detect layers according to LyCORIS detection logic(`weight_list_det`)
-            # https://github.com/KohakuBlueleaf/LyCORIS/tree/8ad8000efb79e2b879054da8c9356e6143591bad/lycoris/modules
-
-            # lora and locon
-            if "lora_up.weight" in values:
-                layer: AnyLoRALayer = LoRALayer(layer_key, values)
-
-            # loha
-            elif "hada_w1_a" in values:
-                layer = LoHALayer(layer_key, values)
-
-            # lokr
-            elif "lokr_w1" in values or "lokr_w1_a" in values:
-                layer = LoKRLayer(layer_key, values)
-
-            # diff
-            elif "diff" in values:
-                layer = FullLayer(layer_key, values)
-
-            # ia3
-            elif "on_input" in values:
-                layer = IA3Layer(layer_key, values)
-
-            # norms
-            elif "w_norm" in values:
-                layer = NormLayer(layer_key, values)
-
-            else:
-                print(f">> Encountered unknown lora layer module in {model.name}: {layer_key} - {list(values.keys())}")
-                raise Exception("Unknown lora format!")
-
-            # lower memory consumption by removing already parsed layer values
-            state_dict[layer_key].clear()
-
-            layer.to(device=device, dtype=dtype)
-            model.layers[layer_key] = layer
-
-        return model
-
-    @staticmethod
-    def _group_state(state_dict: Dict[str, torch.Tensor]) -> Dict[str, Dict[str, torch.Tensor]]:
-        state_dict_groupped: Dict[str, Dict[str, torch.Tensor]] = {}
-
-        for key, value in state_dict.items():
-            stem, leaf = key.split(".", 1)
-            if stem not in state_dict_groupped:
-                state_dict_groupped[stem] = {}
-            state_dict_groupped[stem][leaf] = value
-
-        return state_dict_groupped
-
-
-# code from
-# https://github.com/bmaltais/kohya_ss/blob/2accb1305979ba62f5077a23aabac23b4c37e935/networks/lora_diffusers.py#L15C1-L97C32
-def make_sdxl_unet_conversion_map() -> List[Tuple[str, str]]:
-    """Create a dict mapping state_dict keys from Stability AI SDXL format to diffusers SDXL format."""
-    unet_conversion_map_layer = []
-
-    for i in range(3):  # num_blocks is 3 in sdxl
-        # loop over downblocks/upblocks
-        for j in range(2):
-            # loop over resnets/attentions for downblocks
-            hf_down_res_prefix = f"down_blocks.{i}.resnets.{j}."
-            sd_down_res_prefix = f"input_blocks.{3*i + j + 1}.0."
-            unet_conversion_map_layer.append((sd_down_res_prefix, hf_down_res_prefix))
-
-            if i < 3:
-                # no attention layers in down_blocks.3
-                hf_down_atn_prefix = f"down_blocks.{i}.attentions.{j}."
-                sd_down_atn_prefix = f"input_blocks.{3*i + j + 1}.1."
-                unet_conversion_map_layer.append((sd_down_atn_prefix, hf_down_atn_prefix))
-
-        for j in range(3):
-            # loop over resnets/attentions for upblocks
-            hf_up_res_prefix = f"up_blocks.{i}.resnets.{j}."
-            sd_up_res_prefix = f"output_blocks.{3*i + j}.0."
-            unet_conversion_map_layer.append((sd_up_res_prefix, hf_up_res_prefix))
-
-            # if i > 0: commentout for sdxl
-            # no attention layers in up_blocks.0
-            hf_up_atn_prefix = f"up_blocks.{i}.attentions.{j}."
-            sd_up_atn_prefix = f"output_blocks.{3*i + j}.1."
-            unet_conversion_map_layer.append((sd_up_atn_prefix, hf_up_atn_prefix))
-
-        if i < 3:
-            # no downsample in down_blocks.3
-            hf_downsample_prefix = f"down_blocks.{i}.downsamplers.0.conv."
-            sd_downsample_prefix = f"input_blocks.{3*(i+1)}.0.op."
-            unet_conversion_map_layer.append((sd_downsample_prefix, hf_downsample_prefix))
-
-            # no upsample in up_blocks.3
-            hf_upsample_prefix = f"up_blocks.{i}.upsamplers.0."
-            sd_upsample_prefix = f"output_blocks.{3*i + 2}.{2}."  # change for sdxl
-            unet_conversion_map_layer.append((sd_upsample_prefix, hf_upsample_prefix))
-
-    hf_mid_atn_prefix = "mid_block.attentions.0."
-    sd_mid_atn_prefix = "middle_block.1."
-    unet_conversion_map_layer.append((sd_mid_atn_prefix, hf_mid_atn_prefix))
-
-    for j in range(2):
-        hf_mid_res_prefix = f"mid_block.resnets.{j}."
-        sd_mid_res_prefix = f"middle_block.{2*j}."
-        unet_conversion_map_layer.append((sd_mid_res_prefix, hf_mid_res_prefix))
-
-    unet_conversion_map_resnet = [
-        # (stable-diffusion, HF Diffusers)
-        ("in_layers.0.", "norm1."),
-        ("in_layers.2.", "conv1."),
-        ("out_layers.0.", "norm2."),
-        ("out_layers.3.", "conv2."),
-        ("emb_layers.1.", "time_emb_proj."),
-        ("skip_connection.", "conv_shortcut."),
-    ]
-
-    unet_conversion_map = []
-    for sd, hf in unet_conversion_map_layer:
-        if "resnets" in hf:
-            for sd_res, hf_res in unet_conversion_map_resnet:
-                unet_conversion_map.append((sd + sd_res, hf + hf_res))
-        else:
-            unet_conversion_map.append((sd, hf))
-
-    for j in range(2):
-        hf_time_embed_prefix = f"time_embedding.linear_{j+1}."
-        sd_time_embed_prefix = f"time_embed.{j*2}."
-        unet_conversion_map.append((sd_time_embed_prefix, hf_time_embed_prefix))
-
-    for j in range(2):
-        hf_label_embed_prefix = f"add_embedding.linear_{j+1}."
-        sd_label_embed_prefix = f"label_emb.0.{j*2}."
-        unet_conversion_map.append((sd_label_embed_prefix, hf_label_embed_prefix))
-
-    unet_conversion_map.append(("input_blocks.0.0.", "conv_in."))
-    unet_conversion_map.append(("out.0.", "conv_norm_out."))
-    unet_conversion_map.append(("out.2.", "conv_out."))
-
-    return unet_conversion_map
-
-
-SDXL_UNET_STABILITY_TO_DIFFUSERS_MAP = {
-    sd.rstrip(".").replace(".", "_"): hf.rstrip(".").replace(".", "_") for sd, hf in make_sdxl_unet_conversion_map()
-}

invokeai/backend/lora/conversions/flux_diffusers_lora_conversion_utils.py

@@ -0,0 +1,210 @@
+from typing import Dict
+
+import torch
+
+from invokeai.backend.lora.layers.any_lora_layer import AnyLoRALayer
+from invokeai.backend.lora.layers.concatenated_lora_layer import ConcatenatedLoRALayer
+from invokeai.backend.lora.layers.lora_layer import LoRALayer
+from invokeai.backend.lora.layers.lora_layer_base import LoRALayerBase
+from invokeai.backend.lora.lora_model_raw import LoRAModelRaw
+
+
+def is_state_dict_likely_in_flux_diffusers_format(state_dict: Dict[str, torch.Tensor]) -> bool:
+    """Checks if the provided state dict is likely in the Diffusers FLUX LoRA format.
+
+    This is intended to be a reasonably high-precision detector, but it is not guaranteed to have perfect precision. (A
+    perfect-precision detector would require checking all keys against a whitelist and verifying tensor shapes.)
+    """
+    # First, check that all keys end in "lora_A.weight" or "lora_B.weight" (i.e. are in PEFT format).
+    all_keys_in_peft_format = all(k.endswith(("lora_A.weight", "lora_B.weight")) for k in state_dict.keys())
+
+    # Next, check that this is likely a FLUX model by spot-checking a few keys.
+    expected_keys = [
+        "transformer.single_transformer_blocks.0.attn.to_q.lora_A.weight",
+        "transformer.single_transformer_blocks.0.attn.to_q.lora_B.weight",
+        "transformer.transformer_blocks.0.attn.add_q_proj.lora_A.weight",
+        "transformer.transformer_blocks.0.attn.add_q_proj.lora_B.weight",
+    ]
+    all_expected_keys_present = all(k in state_dict for k in expected_keys)
+
+    return all_keys_in_peft_format and all_expected_keys_present
+
+
+# TODO(ryand): What alpha should we use? 1.0? Rank of the LoRA?
+def lora_model_from_flux_diffusers_state_dict(state_dict: Dict[str, torch.Tensor], alpha: float = 1.0) -> LoRAModelRaw:  # pyright: ignore[reportRedeclaration] (state_dict is intentionally re-declared)
+    """Loads a state dict in the Diffusers FLUX LoRA format into a LoRAModelRaw object.
+
+    This function is based on:
+    https://github.com/huggingface/diffusers/blob/55ac421f7bb12fd00ccbef727be4dc2f3f920abb/scripts/convert_flux_to_diffusers.py
+    """
+    # Group keys by layer.
+    grouped_state_dict: dict[str, dict[str, torch.Tensor]] = _group_by_layer(state_dict)
+
+    # Remove the "transformer." prefix from all keys.
+    grouped_state_dict = {k.replace("transformer.", ""): v for k, v in grouped_state_dict.items()}
+
+    # Constants for FLUX.1
+    num_double_layers = 19
+    num_single_layers = 38
+    # inner_dim = 3072
+    # mlp_ratio = 4.0
+
+    layers: dict[str, AnyLoRALayer] = {}
+
+    def add_lora_layer_if_present(src_key: str, dst_key: str) -> None:
+        if src_key in grouped_state_dict:
+            src_layer_dict = grouped_state_dict.pop(src_key)
+            layers[dst_key] = LoRALayer(
+                values={
+                    "lora_down.weight": src_layer_dict.pop("lora_A.weight"),
+                    "lora_up.weight": src_layer_dict.pop("lora_B.weight"),
+                    "alpha": torch.tensor(alpha),
+                },
+            )
+            assert len(src_layer_dict) == 0
+
+    def add_qkv_lora_layer_if_present(src_keys: list[str], dst_qkv_key: str) -> None:
+        """Handle the Q, K, V matrices for a transformer block. We need special handling because the diffusers format
+        stores them in separate matrices, whereas the BFL format used internally by InvokeAI concatenates them.
+        """
+        # We expect that either all src keys are present or none of them are. Verify this.
+        keys_present = [key in grouped_state_dict for key in src_keys]
+        assert all(keys_present) or not any(keys_present)
+
+        # If none of the keys are present, return early.
+        if not any(keys_present):
+            return
+
+        src_layer_dicts = [grouped_state_dict.pop(key) for key in src_keys]
+        sub_layers: list[LoRALayerBase] = []
+        for src_layer_dict in src_layer_dicts:
+            sub_layers.append(
+                LoRALayer(
+                    values={
+                        "lora_down.weight": src_layer_dict.pop("lora_A.weight"),
+                        "lora_up.weight": src_layer_dict.pop("lora_B.weight"),
+                        "alpha": torch.tensor(alpha),
+                    },
+                )
+            )
+            assert len(src_layer_dict) == 0
+        layers[dst_qkv_key] = ConcatenatedLoRALayer(lora_layers=sub_layers, concat_axis=0)
+
+    # time_text_embed.timestep_embedder -> time_in.
+    add_lora_layer_if_present("time_text_embed.timestep_embedder.linear_1", "time_in.in_layer")
+    add_lora_layer_if_present("time_text_embed.timestep_embedder.linear_2", "time_in.out_layer")
+
+    # time_text_embed.text_embedder -> vector_in.
+    add_lora_layer_if_present("time_text_embed.text_embedder.linear_1", "vector_in.in_layer")
+    add_lora_layer_if_present("time_text_embed.text_embedder.linear_2", "vector_in.out_layer")
+
+    # time_text_embed.guidance_embedder -> guidance_in.
+    add_lora_layer_if_present("time_text_embed.guidance_embedder.linear_1", "guidance_in")
+    add_lora_layer_if_present("time_text_embed.guidance_embedder.linear_2", "guidance_in")
+
+    # context_embedder -> txt_in.
+    add_lora_layer_if_present("context_embedder", "txt_in")
+
+    # x_embedder -> img_in.
+    add_lora_layer_if_present("x_embedder", "img_in")
+
+    # Double transformer blocks.
+    for i in range(num_double_layers):
+        # norms.
+        add_lora_layer_if_present(f"transformer_blocks.{i}.norm1.linear", f"double_blocks.{i}.img_mod.lin")
+        add_lora_layer_if_present(f"transformer_blocks.{i}.norm1_context.linear", f"double_blocks.{i}.txt_mod.lin")
+
+        # Q, K, V
+        add_qkv_lora_layer_if_present(
+            [
+                f"transformer_blocks.{i}.attn.to_q",
+                f"transformer_blocks.{i}.attn.to_k",
+                f"transformer_blocks.{i}.attn.to_v",
+            ],
+            f"double_blocks.{i}.img_attn.qkv",
+        )
+        add_qkv_lora_layer_if_present(
+            [
+                f"transformer_blocks.{i}.attn.add_q_proj",
+                f"transformer_blocks.{i}.attn.add_k_proj",
+                f"transformer_blocks.{i}.attn.add_v_proj",
+            ],
+            f"double_blocks.{i}.txt_attn.qkv",
+        )
+
+        # ff img_mlp
+        add_lora_layer_if_present(
+            f"transformer_blocks.{i}.ff.net.0.proj",
+            f"double_blocks.{i}.img_mlp.0",
+        )
+        add_lora_layer_if_present(
+            f"transformer_blocks.{i}.ff.net.2",
+            f"double_blocks.{i}.img_mlp.2",
+        )
+
+        # ff txt_mlp
+        add_lora_layer_if_present(
+            f"transformer_blocks.{i}.ff_context.net.0.proj",
+            f"double_blocks.{i}.txt_mlp.0",
+        )
+        add_lora_layer_if_present(
+            f"transformer_blocks.{i}.ff_context.net.2",
+            f"double_blocks.{i}.txt_mlp.2",
+        )
+
+        # output projections.
+        add_lora_layer_if_present(
+            f"transformer_blocks.{i}.attn.to_out.0",
+            f"double_blocks.{i}.img_attn.proj",
+        )
+        add_lora_layer_if_present(
+            f"transformer_blocks.{i}.attn.to_add_out",
+            f"double_blocks.{i}.txt_attn.proj",
+        )
+
+    # Single transformer blocks.
+    for i in range(num_single_layers):
+        # norms
+        add_lora_layer_if_present(
+            f"single_transformer_blocks.{i}.norm.linear",
+            f"single_blocks.{i}.modulation.lin",
+        )
+
+        # Q, K, V, mlp
+        add_qkv_lora_layer_if_present(
+            [
+                f"single_transformer_blocks.{i}.attn.to_q",
+                f"single_transformer_blocks.{i}.attn.to_k",
+                f"single_transformer_blocks.{i}.attn.to_v",
+                f"single_transformer_blocks.{i}.proj_mlp",
+            ],
+            f"single_blocks.{i}.linear1",
+        )
+
+        # Output projections.
+        add_lora_layer_if_present(
+            f"single_transformer_blocks.{i}.proj_out",
+            f"single_blocks.{i}.linear2",
+        )
+
+    # Final layer.
+    add_lora_layer_if_present("proj_out", "final_layer.linear")
+
+    # Assert that all keys were processed.
+    assert len(grouped_state_dict) == 0
+
+    return LoRAModelRaw(layers=layers)
+
+
+def _group_by_layer(state_dict: Dict[str, torch.Tensor]) -> dict[str, dict[str, torch.Tensor]]:
+    """Groups the keys in the state dict by layer."""
+    layer_dict: dict[str, dict[str, torch.Tensor]] = {}
+    for key in state_dict:
+        # Split the 'lora_A.weight' or 'lora_B.weight' suffix from the layer name.
+        parts = key.rsplit(".", maxsplit=2)
+        layer_name = parts[0]
+        key_name = ".".join(parts[1:])
+        if layer_name not in layer_dict:
+            layer_dict[layer_name] = {}
+        layer_dict[layer_name][key_name] = state_dict[key]
+    return layer_dict

invokeai/backend/lora/conversions/flux_kohya_lora_conversion_utils.py

@@ -0,0 +1,80 @@
+import re
+from typing import Any, Dict, TypeVar
+
+import torch
+
+from invokeai.backend.lora.layers.any_lora_layer import AnyLoRALayer
+from invokeai.backend.lora.layers.utils import any_lora_layer_from_state_dict
+from invokeai.backend.lora.lora_model_raw import LoRAModelRaw
+
+# A regex pattern that matches all of the keys in the Kohya FLUX LoRA format.
+# Example keys:
+#   lora_unet_double_blocks_0_img_attn_proj.alpha
+#   lora_unet_double_blocks_0_img_attn_proj.lora_down.weight
+#   lora_unet_double_blocks_0_img_attn_proj.lora_up.weight
+FLUX_KOHYA_KEY_REGEX = (
+    r"lora_unet_(\w+_blocks)_(\d+)_(img_attn|img_mlp|img_mod|txt_attn|txt_mlp|txt_mod|linear1|linear2|modulation)_?(.*)"
+)
+
+
+def is_state_dict_likely_in_flux_kohya_format(state_dict: Dict[str, Any]) -> bool:
+    """Checks if the provided state dict is likely in the Kohya FLUX LoRA format.
+
+    This is intended to be a high-precision detector, but it is not guaranteed to have perfect precision. (A
+    perfect-precision detector would require checking all keys against a whitelist and verifying tensor shapes.)
+    """
+    return all(re.match(FLUX_KOHYA_KEY_REGEX, k) for k in state_dict.keys())
+
+
+def lora_model_from_flux_kohya_state_dict(state_dict: Dict[str, torch.Tensor]) -> LoRAModelRaw:
+    # Group keys by layer.
+    grouped_state_dict: dict[str, dict[str, torch.Tensor]] = {}
+    for key, value in state_dict.items():
+        layer_name, param_name = key.split(".", 1)
+        if layer_name not in grouped_state_dict:
+            grouped_state_dict[layer_name] = {}
+        grouped_state_dict[layer_name][param_name] = value
+
+    # Convert the state dict to the InvokeAI format.
+    grouped_state_dict = convert_flux_kohya_state_dict_to_invoke_format(grouped_state_dict)
+
+    # Create LoRA layers.
+    layers: dict[str, AnyLoRALayer] = {}
+    for layer_key, layer_state_dict in grouped_state_dict.items():
+        layers[layer_key] = any_lora_layer_from_state_dict(layer_state_dict)
+
+    # Create and return the LoRAModelRaw.
+    return LoRAModelRaw(layers=layers)
+
+
+T = TypeVar("T")
+
+
+def convert_flux_kohya_state_dict_to_invoke_format(state_dict: Dict[str, T]) -> Dict[str, T]:
+    """Converts a state dict from the Kohya FLUX LoRA format to LoRA weight format used internally by InvokeAI.
+
+    Example key conversions:
+    "lora_unet_double_blocks_0_img_attn_proj" -> "double_blocks.0.img_attn.proj"
+    "lora_unet_double_blocks_0_img_attn_proj" -> "double_blocks.0.img_attn.proj"
+    "lora_unet_double_blocks_0_img_attn_proj" -> "double_blocks.0.img_attn.proj"
+    "lora_unet_double_blocks_0_img_attn_qkv" -> "double_blocks.0.img_attn.qkv"
+    "lora_unet_double_blocks_0_img_attn_qkv" -> "double_blocks.0.img.attn.qkv"
+    "lora_unet_double_blocks_0_img_attn_qkv" -> "double_blocks.0.img.attn.qkv"
+    """
+
+    def replace_func(match: re.Match[str]) -> str:
+        s = f"{match.group(1)}.{match.group(2)}.{match.group(3)}"
+        if match.group(4):
+            s += f".{match.group(4)}"
+        return s
+
+    converted_dict: dict[str, T] = {}
+    for k, v in state_dict.items():
+        match = re.match(FLUX_KOHYA_KEY_REGEX, k)
+        if match:
+            new_key = re.sub(FLUX_KOHYA_KEY_REGEX, replace_func, k)
+            converted_dict[new_key] = v
+        else:
+            raise ValueError(f"Key '{k}' does not match the expected pattern for FLUX LoRA weights.")
+
+    return converted_dict

invokeai/backend/lora/conversions/sd_lora_conversion_utils.py

@@ -0,0 +1,29 @@
+from typing import Dict
+
+import torch
+
+from invokeai.backend.lora.layers.any_lora_layer import AnyLoRALayer
+from invokeai.backend.lora.layers.utils import any_lora_layer_from_state_dict
+from invokeai.backend.lora.lora_model_raw import LoRAModelRaw
+
+
+def lora_model_from_sd_state_dict(state_dict: Dict[str, torch.Tensor]) -> LoRAModelRaw:
+    grouped_state_dict: dict[str, dict[str, torch.Tensor]] = _group_state(state_dict)
+
+    layers: dict[str, AnyLoRALayer] = {}
+    for layer_key, values in grouped_state_dict.items():
+        layers[layer_key] = any_lora_layer_from_state_dict(values)
+
+    return LoRAModelRaw(layers=layers)
+
+
+def _group_state(state_dict: Dict[str, torch.Tensor]) -> Dict[str, Dict[str, torch.Tensor]]:
+    state_dict_groupped: Dict[str, Dict[str, torch.Tensor]] = {}
+
+    for key, value in state_dict.items():
+        stem, leaf = key.split(".", 1)
+        if stem not in state_dict_groupped:
+            state_dict_groupped[stem] = {}
+        state_dict_groupped[stem][leaf] = value
+
+    return state_dict_groupped

invokeai/backend/lora/conversions/sdxl_lora_conversion_utils.py

@@ -0,0 +1,154 @@
+import bisect
+from typing import Dict, List, Tuple, TypeVar
+
+T = TypeVar("T")
+
+
+def convert_sdxl_keys_to_diffusers_format(state_dict: Dict[str, T]) -> dict[str, T]:
+    """Convert the keys of an SDXL LoRA state_dict to diffusers format.
+
+    The input state_dict can be in either Stability AI format or diffusers format. If the state_dict is already in
+    diffusers format, then this function will have no effect.
+
+    This function is adapted from:
+    https://github.com/bmaltais/kohya_ss/blob/2accb1305979ba62f5077a23aabac23b4c37e935/networks/lora_diffusers.py#L385-L409
+
+    Args:
+        state_dict (Dict[str, Tensor]): The SDXL LoRA state_dict.
+
+    Raises:
+        ValueError: If state_dict contains an unrecognized key, or not all keys could be converted.
+
+    Returns:
+        Dict[str, Tensor]: The diffusers-format state_dict.
+    """
+    converted_count = 0  # The number of Stability AI keys converted to diffusers format.
+    not_converted_count = 0  # The number of keys that were not converted.
+
+    # Get a sorted list of Stability AI UNet keys so that we can efficiently search for keys with matching prefixes.
+    # For example, we want to efficiently find `input_blocks_4_1` in the list when searching for
+    # `input_blocks_4_1_proj_in`.
+    stability_unet_keys = list(SDXL_UNET_STABILITY_TO_DIFFUSERS_MAP)
+    stability_unet_keys.sort()
+
+    new_state_dict: dict[str, T] = {}
+    for full_key, value in state_dict.items():
+        if full_key.startswith("lora_unet_"):
+            search_key = full_key.replace("lora_unet_", "")
+            # Use bisect to find the key in stability_unet_keys that *may* match the search_key's prefix.
+            position = bisect.bisect_right(stability_unet_keys, search_key)
+            map_key = stability_unet_keys[position - 1]
+            # Now, check if the map_key *actually* matches the search_key.
+            if search_key.startswith(map_key):
+                new_key = full_key.replace(map_key, SDXL_UNET_STABILITY_TO_DIFFUSERS_MAP[map_key])
+                new_state_dict[new_key] = value
+                converted_count += 1
+            else:
+                new_state_dict[full_key] = value
+                not_converted_count += 1
+        elif full_key.startswith("lora_te1_") or full_key.startswith("lora_te2_"):
+            # The CLIP text encoders have the same keys in both Stability AI and diffusers formats.
+            new_state_dict[full_key] = value
+            continue
+        else:
+            raise ValueError(f"Unrecognized SDXL LoRA key prefix: '{full_key}'.")
+
+    if converted_count > 0 and not_converted_count > 0:
+        raise ValueError(
+            f"The SDXL LoRA could only be partially converted to diffusers format. converted={converted_count},"
+            f" not_converted={not_converted_count}"
+        )
+
+    return new_state_dict
+
+
+# code from
+# https://github.com/bmaltais/kohya_ss/blob/2accb1305979ba62f5077a23aabac23b4c37e935/networks/lora_diffusers.py#L15C1-L97C32
+def _make_sdxl_unet_conversion_map() -> List[Tuple[str, str]]:
+    """Create a dict mapping state_dict keys from Stability AI SDXL format to diffusers SDXL format."""
+    unet_conversion_map_layer: list[tuple[str, str]] = []
+
+    for i in range(3):  # num_blocks is 3 in sdxl
+        # loop over downblocks/upblocks
+        for j in range(2):
+            # loop over resnets/attentions for downblocks
+            hf_down_res_prefix = f"down_blocks.{i}.resnets.{j}."
+            sd_down_res_prefix = f"input_blocks.{3*i + j + 1}.0."
+            unet_conversion_map_layer.append((sd_down_res_prefix, hf_down_res_prefix))
+
+            if i < 3:
+                # no attention layers in down_blocks.3
+                hf_down_atn_prefix = f"down_blocks.{i}.attentions.{j}."
+                sd_down_atn_prefix = f"input_blocks.{3*i + j + 1}.1."
+                unet_conversion_map_layer.append((sd_down_atn_prefix, hf_down_atn_prefix))
+
+        for j in range(3):
+            # loop over resnets/attentions for upblocks
+            hf_up_res_prefix = f"up_blocks.{i}.resnets.{j}."
+            sd_up_res_prefix = f"output_blocks.{3*i + j}.0."
+            unet_conversion_map_layer.append((sd_up_res_prefix, hf_up_res_prefix))
+
+            # if i > 0: commentout for sdxl
+            # no attention layers in up_blocks.0
+            hf_up_atn_prefix = f"up_blocks.{i}.attentions.{j}."
+            sd_up_atn_prefix = f"output_blocks.{3*i + j}.1."
+            unet_conversion_map_layer.append((sd_up_atn_prefix, hf_up_atn_prefix))
+
+        if i < 3:
+            # no downsample in down_blocks.3
+            hf_downsample_prefix = f"down_blocks.{i}.downsamplers.0.conv."
+            sd_downsample_prefix = f"input_blocks.{3*(i+1)}.0.op."
+            unet_conversion_map_layer.append((sd_downsample_prefix, hf_downsample_prefix))
+
+            # no upsample in up_blocks.3
+            hf_upsample_prefix = f"up_blocks.{i}.upsamplers.0."
+            sd_upsample_prefix = f"output_blocks.{3*i + 2}.{2}."  # change for sdxl
+            unet_conversion_map_layer.append((sd_upsample_prefix, hf_upsample_prefix))
+
+    hf_mid_atn_prefix = "mid_block.attentions.0."
+    sd_mid_atn_prefix = "middle_block.1."
+    unet_conversion_map_layer.append((sd_mid_atn_prefix, hf_mid_atn_prefix))
+
+    for j in range(2):
+        hf_mid_res_prefix = f"mid_block.resnets.{j}."
+        sd_mid_res_prefix = f"middle_block.{2*j}."
+        unet_conversion_map_layer.append((sd_mid_res_prefix, hf_mid_res_prefix))
+
+    unet_conversion_map_resnet = [
+        # (stable-diffusion, HF Diffusers)
+        ("in_layers.0.", "norm1."),
+        ("in_layers.2.", "conv1."),
+        ("out_layers.0.", "norm2."),
+        ("out_layers.3.", "conv2."),
+        ("emb_layers.1.", "time_emb_proj."),
+        ("skip_connection.", "conv_shortcut."),
+    ]
+
+    unet_conversion_map: list[tuple[str, str]] = []
+    for sd, hf in unet_conversion_map_layer:
+        if "resnets" in hf:
+            for sd_res, hf_res in unet_conversion_map_resnet:
+                unet_conversion_map.append((sd + sd_res, hf + hf_res))
+        else:
+            unet_conversion_map.append((sd, hf))
+
+    for j in range(2):
+        hf_time_embed_prefix = f"time_embedding.linear_{j+1}."
+        sd_time_embed_prefix = f"time_embed.{j*2}."
+        unet_conversion_map.append((sd_time_embed_prefix, hf_time_embed_prefix))
+
+    for j in range(2):
+        hf_label_embed_prefix = f"add_embedding.linear_{j+1}."
+        sd_label_embed_prefix = f"label_emb.0.{j*2}."
+        unet_conversion_map.append((sd_label_embed_prefix, hf_label_embed_prefix))
+
+    unet_conversion_map.append(("input_blocks.0.0.", "conv_in."))
+    unet_conversion_map.append(("out.0.", "conv_norm_out."))
+    unet_conversion_map.append(("out.2.", "conv_out."))
+
+    return unet_conversion_map
+
+
+SDXL_UNET_STABILITY_TO_DIFFUSERS_MAP = {
+    sd.rstrip(".").replace(".", "_"): hf.rstrip(".").replace(".", "_") for sd, hf in _make_sdxl_unet_conversion_map()
+}

invokeai/backend/lora/layers/any_lora_layer.py

@@ -0,0 +1,11 @@
+from typing import Union
+
+from invokeai.backend.lora.layers.concatenated_lora_layer import ConcatenatedLoRALayer
+from invokeai.backend.lora.layers.full_layer import FullLayer
+from invokeai.backend.lora.layers.ia3_layer import IA3Layer
+from invokeai.backend.lora.layers.loha_layer import LoHALayer
+from invokeai.backend.lora.layers.lokr_layer import LoKRLayer
+from invokeai.backend.lora.layers.lora_layer import LoRALayer
+from invokeai.backend.lora.layers.norm_layer import NormLayer
+
+AnyLoRALayer = Union[LoRALayer, LoHALayer, LoKRLayer, FullLayer, IA3Layer, NormLayer, ConcatenatedLoRALayer]

invokeai/backend/lora/layers/concatenated_lora_layer.py

@@ -0,0 +1,46 @@
+from typing import List, Optional
+
+import torch
+
+from invokeai.backend.lora.layers.lora_layer_base import LoRALayerBase
+
+
+class ConcatenatedLoRALayer(LoRALayerBase):
+    """A LoRA layer that is composed of multiple LoRA layers concatenated along a specified axis.
+
+    This class was created to handle a special case with FLUX LoRA models. In the BFL FLUX model format, the attention
+    Q, K, V matrices are concatenated along the first dimension. In the diffusers LoRA format, the Q, K, V matrices are
+    stored as separate tensors. This class enables diffusers LoRA layers to be used in BFL FLUX models.
+    """
+
+    def __init__(self, lora_layers: List[LoRALayerBase], concat_axis: int = 0):
+        # Note: We pass values={} to the base class, because the values are handled by the individual LoRA layers.
+        super().__init__(values={})
+
+        self._lora_layers = lora_layers
+        self._concat_axis = concat_axis
+
+    def get_weight(self, orig_weight: torch.Tensor) -> torch.Tensor:
+        # TODO(ryand): Currently, we pass orig_weight=None to the sub-layers. If we want to support sub-layers that
+        # require this value, we will need to implement chunking of the original weight tensor here.
+        layer_weights = [lora_layer.get_weight(None) for lora_layer in self._lora_layers]  # pyright: ignore[reportArgumentType]
+        return torch.cat(layer_weights, dim=self._concat_axis)
+
+    def get_bias(self, orig_bias: torch.Tensor) -> Optional[torch.Tensor]:
+        # TODO(ryand): Currently, we pass orig_bias=None to the sub-layers. If we want to support sub-layers that
+        # require this value, we will need to implement chunking of the original bias tensor here.
+        layer_biases = [lora_layer.get_bias(None) for lora_layer in self._lora_layers]  # pyright: ignore[reportArgumentType]
+        layer_bias_is_none = [layer_bias is None for layer_bias in layer_biases]
+        if any(layer_bias_is_none):
+            assert all(layer_bias_is_none)
+            return None
+
+        # Ignore the type error, because we have just verified that all layer biases are non-None.
+        return torch.cat(layer_biases, dim=self._concat_axis)
+
+    def calc_size(self) -> int:
+        return sum(lora_layer.calc_size() for lora_layer in self._lora_layers)
+
+    def to(self, device: Optional[torch.device] = None, dtype: Optional[torch.dtype] = None) -> None:
+        for lora_layer in self._lora_layers:
+            lora_layer.to(device=device, dtype=dtype)

invokeai/backend/lora/layers/full_layer.py

@@ -0,0 +1,36 @@
+from typing import Dict, Optional
+
+import torch
+
+from invokeai.backend.lora.layers.lora_layer_base import LoRALayerBase
+
+
+class FullLayer(LoRALayerBase):
+    # bias handled in LoRALayerBase(calc_size, to)
+    # weight: torch.Tensor
+    # bias: Optional[torch.Tensor]
+
+    def __init__(
+        self,
+        values: Dict[str, torch.Tensor],
+    ):
+        super().__init__(values)
+
+        self.weight = values["diff"]
+        self.bias = values.get("diff_b", None)
+
+        self.rank = None  # unscaled
+        self.check_keys(values, {"diff", "diff_b"})
+
+    def get_weight(self, orig_weight: torch.Tensor) -> torch.Tensor:
+        return self.weight
+
+    def calc_size(self) -> int:
+        model_size = super().calc_size()
+        model_size += self.weight.nelement() * self.weight.element_size()
+        return model_size
+
+    def to(self, device: Optional[torch.device] = None, dtype: Optional[torch.dtype] = None) -> None:
+        super().to(device=device, dtype=dtype)
+
+        self.weight = self.weight.to(device=device, dtype=dtype)

invokeai/backend/lora/layers/ia3_layer.py

@@ -0,0 +1,41 @@
+from typing import Dict, Optional
+
+import torch
+
+from invokeai.backend.lora.layers.lora_layer_base import LoRALayerBase
+
+
+class IA3Layer(LoRALayerBase):
+    # weight: torch.Tensor
+    # on_input: torch.Tensor
+
+    def __init__(
+        self,
+        values: Dict[str, torch.Tensor],
+    ):
+        super().__init__(values)
+
+        self.weight = values["weight"]
+        self.on_input = values["on_input"]
+
+        self.rank = None  # unscaled
+        self.check_keys(values, {"weight", "on_input"})
+
+    def get_weight(self, orig_weight: torch.Tensor) -> torch.Tensor:
+        weight = self.weight
+        if not self.on_input:
+            weight = weight.reshape(-1, 1)
+        assert orig_weight is not None
+        return orig_weight * weight
+
+    def calc_size(self) -> int:
+        model_size = super().calc_size()
+        model_size += self.weight.nelement() * self.weight.element_size()
+        model_size += self.on_input.nelement() * self.on_input.element_size()
+        return model_size
+
+    def to(self, device: Optional[torch.device] = None, dtype: Optional[torch.dtype] = None):
+        super().to(device=device, dtype=dtype)
+
+        self.weight = self.weight.to(device=device, dtype=dtype)
+        self.on_input = self.on_input.to(device=device, dtype=dtype)

invokeai/backend/lora/layers/loha_layer.py

@@ -0,0 +1,68 @@
+from typing import Dict, Optional
+
+import torch
+
+from invokeai.backend.lora.layers.lora_layer_base import LoRALayerBase
+
+
+class LoHALayer(LoRALayerBase):
+    # w1_a: torch.Tensor
+    # w1_b: torch.Tensor
+    # w2_a: torch.Tensor
+    # w2_b: torch.Tensor
+    # t1: Optional[torch.Tensor] = None
+    # t2: Optional[torch.Tensor] = None
+
+    def __init__(self, values: Dict[str, torch.Tensor]):
+        super().__init__(values)
+
+        self.w1_a = values["hada_w1_a"]
+        self.w1_b = values["hada_w1_b"]
+        self.w2_a = values["hada_w2_a"]
+        self.w2_b = values["hada_w2_b"]
+        self.t1 = values.get("hada_t1", None)
+        self.t2 = values.get("hada_t2", None)
+
+        self.rank = self.w1_b.shape[0]
+        self.check_keys(
+            values,
+            {
+                "hada_w1_a",
+                "hada_w1_b",
+                "hada_w2_a",
+                "hada_w2_b",
+                "hada_t1",
+                "hada_t2",
+            },
+        )
+
+    def get_weight(self, orig_weight: torch.Tensor) -> torch.Tensor:
+        if self.t1 is None:
+            weight: torch.Tensor = (self.w1_a @ self.w1_b) * (self.w2_a @ self.w2_b)
+
+        else:
+            rebuild1 = torch.einsum("i j k l, j r, i p -> p r k l", self.t1, self.w1_b, self.w1_a)
+            rebuild2 = torch.einsum("i j k l, j r, i p -> p r k l", self.t2, self.w2_b, self.w2_a)
+            weight = rebuild1 * rebuild2
+
+        return weight
+
+    def calc_size(self) -> int:
+        model_size = super().calc_size()
+        for val in [self.w1_a, self.w1_b, self.w2_a, self.w2_b, self.t1, self.t2]:
+            if val is not None:
+                model_size += val.nelement() * val.element_size()
+        return model_size
+
+    def to(self, device: Optional[torch.device] = None, dtype: Optional[torch.dtype] = None) -> None:
+        super().to(device=device, dtype=dtype)
+
+        self.w1_a = self.w1_a.to(device=device, dtype=dtype)
+        self.w1_b = self.w1_b.to(device=device, dtype=dtype)
+        if self.t1 is not None:
+            self.t1 = self.t1.to(device=device, dtype=dtype)
+
+        self.w2_a = self.w2_a.to(device=device, dtype=dtype)
+        self.w2_b = self.w2_b.to(device=device, dtype=dtype)
+        if self.t2 is not None:
+            self.t2 = self.t2.to(device=device, dtype=dtype)

invokeai/backend/lora/layers/lokr_layer.py

@@ -0,0 +1,113 @@
+from typing import Dict, Optional
+
+import torch
+
+from invokeai.backend.lora.layers.lora_layer_base import LoRALayerBase
+
+
+class LoKRLayer(LoRALayerBase):
+    # w1: Optional[torch.Tensor] = None
+    # w1_a: Optional[torch.Tensor] = None
+    # w1_b: Optional[torch.Tensor] = None
+    # w2: Optional[torch.Tensor] = None
+    # w2_a: Optional[torch.Tensor] = None
+    # w2_b: Optional[torch.Tensor] = None
+    # t2: Optional[torch.Tensor] = None
+
+    def __init__(
+        self,
+        values: Dict[str, torch.Tensor],
+    ):
+        super().__init__(values)
+
+        self.w1 = values.get("lokr_w1", None)
+        if self.w1 is None:
+            self.w1_a = values["lokr_w1_a"]
+            self.w1_b = values["lokr_w1_b"]
+        else:
+            self.w1_b = None
+            self.w1_a = None
+
+        self.w2 = values.get("lokr_w2", None)
+        if self.w2 is None:
+            self.w2_a = values["lokr_w2_a"]
+            self.w2_b = values["lokr_w2_b"]
+        else:
+            self.w2_a = None
+            self.w2_b = None
+
+        self.t2 = values.get("lokr_t2", None)
+
+        if self.w1_b is not None:
+            self.rank = self.w1_b.shape[0]
+        elif self.w2_b is not None:
+            self.rank = self.w2_b.shape[0]
+        else:
+            self.rank = None  # unscaled
+
+        self.check_keys(
+            values,
+            {
+                "lokr_w1",
+                "lokr_w1_a",
+                "lokr_w1_b",
+                "lokr_w2",
+                "lokr_w2_a",
+                "lokr_w2_b",
+                "lokr_t2",
+            },
+        )
+
+    def get_weight(self, orig_weight: torch.Tensor) -> torch.Tensor:
+        w1: Optional[torch.Tensor] = self.w1
+        if w1 is None:
+            assert self.w1_a is not None
+            assert self.w1_b is not None
+            w1 = self.w1_a @ self.w1_b
+
+        w2 = self.w2
+        if w2 is None:
+            if self.t2 is None:
+                assert self.w2_a is not None
+                assert self.w2_b is not None
+                w2 = self.w2_a @ self.w2_b
+            else:
+                w2 = torch.einsum("i j k l, i p, j r -> p r k l", self.t2, self.w2_a, self.w2_b)
+
+        if len(w2.shape) == 4:
+            w1 = w1.unsqueeze(2).unsqueeze(2)
+        w2 = w2.contiguous()
+        assert w1 is not None
+        assert w2 is not None
+        weight = torch.kron(w1, w2)
+
+        return weight
+
+    def calc_size(self) -> int:
+        model_size = super().calc_size()
+        for val in [self.w1, self.w1_a, self.w1_b, self.w2, self.w2_a, self.w2_b, self.t2]:
+            if val is not None:
+                model_size += val.nelement() * val.element_size()
+        return model_size
+
+    def to(self, device: Optional[torch.device] = None, dtype: Optional[torch.dtype] = None) -> None:
+        super().to(device=device, dtype=dtype)
+
+        if self.w1 is not None:
+            self.w1 = self.w1.to(device=device, dtype=dtype)
+        else:
+            assert self.w1_a is not None
+            assert self.w1_b is not None
+            self.w1_a = self.w1_a.to(device=device, dtype=dtype)
+            self.w1_b = self.w1_b.to(device=device, dtype=dtype)
+
+        if self.w2 is not None:
+            self.w2 = self.w2.to(device=device, dtype=dtype)
+        else:
+            assert self.w2_a is not None
+            assert self.w2_b is not None
+            self.w2_a = self.w2_a.to(device=device, dtype=dtype)
+            self.w2_b = self.w2_b.to(device=device, dtype=dtype)
+
+        if self.t2 is not None:
+            self.t2 = self.t2.to(device=device, dtype=dtype)

invokeai/backend/lora/layers/lora_layer.py

@@ -0,0 +1,58 @@
+from typing import Dict, Optional
+
+import torch
+
+from invokeai.backend.lora.layers.lora_layer_base import LoRALayerBase
+
+
+# TODO: find and debug lora/locon with bias
+class LoRALayer(LoRALayerBase):
+    # up: torch.Tensor
+    # mid: Optional[torch.Tensor]
+    # down: torch.Tensor
+
+    def __init__(
+        self,
+        values: Dict[str, torch.Tensor],
+    ):
+        super().__init__(values)
+
+        self.up = values["lora_up.weight"]
+        self.down = values["lora_down.weight"]
+        self.mid = values.get("lora_mid.weight", None)
+
+        self.rank = self.down.shape[0]
+        self.check_keys(
+            values,
+            {
+                "lora_up.weight",
+                "lora_down.weight",
+                "lora_mid.weight",
+            },
+        )
+
+    def get_weight(self, orig_weight: torch.Tensor) -> torch.Tensor:
+        if self.mid is not None:
+            up = self.up.reshape(self.up.shape[0], self.up.shape[1])
+            down = self.down.reshape(self.down.shape[0], self.down.shape[1])
+            weight = torch.einsum("m n w h, i m, n j -> i j w h", self.mid, up, down)
+        else:
+            weight = self.up.reshape(self.up.shape[0], -1) @ self.down.reshape(self.down.shape[0], -1)
+
+        return weight
+
+    def calc_size(self) -> int:
+        model_size = super().calc_size()
+        for val in [self.up, self.mid, self.down]:
+            if val is not None:
+                model_size += val.nelement() * val.element_size()
+        return model_size
+
+    def to(self, device: Optional[torch.device] = None, dtype: Optional[torch.dtype] = None) -> None:
+        super().to(device=device, dtype=dtype)
+
+        self.up = self.up.to(device=device, dtype=dtype)
+        self.down = self.down.to(device=device, dtype=dtype)
+
+        if self.mid is not None:
+            self.mid = self.mid.to(device=device, dtype=dtype)

invokeai/backend/lora/layers/lora_layer_base.py

@@ -0,0 +1,71 @@
+from typing import Dict, Optional, Set
+
+import torch
+
+import invokeai.backend.util.logging as logger
+
+
+class LoRALayerBase:
+    # rank: Optional[int]
+    # alpha: Optional[float]
+    # bias: Optional[torch.Tensor]
+
+    # @property
+    # def scale(self):
+    #    return self.alpha / self.rank if (self.alpha and self.rank) else 1.0
+
+    def __init__(
+        self,
+        values: Dict[str, torch.Tensor],
+    ):
+        if "alpha" in values:
+            self.alpha = values["alpha"].item()
+        else:
+            self.alpha = None
+
+        if "bias_indices" in values and "bias_values" in values and "bias_size" in values:
+            self.bias: Optional[torch.Tensor] = torch.sparse_coo_tensor(
+                values["bias_indices"],
+                values["bias_values"],
+                tuple(values["bias_size"]),
+            )
+
+        else:
+            self.bias = None
+
+        self.rank = None  # set in layer implementation
+
+    def get_weight(self, orig_weight: torch.Tensor) -> torch.Tensor:
+        raise NotImplementedError()
+
+    def get_bias(self, orig_bias: torch.Tensor) -> Optional[torch.Tensor]:
+        return self.bias
+
+    def get_parameters(self, orig_module: torch.nn.Module) -> Dict[str, torch.Tensor]:
+        params = {"weight": self.get_weight(orig_module.weight)}
+        bias = self.get_bias(orig_module.bias)
+        if bias is not None:
+            params["bias"] = bias
+        return params
+
+    def calc_size(self) -> int:
+        model_size = 0
+        for val in [self.bias]:
+            if val is not None:
+                model_size += val.nelement() * val.element_size()
+        return model_size
+
+    def to(self, device: Optional[torch.device] = None, dtype: Optional[torch.dtype] = None) -> None:
+        if self.bias is not None:
+            self.bias = self.bias.to(device=device, dtype=dtype)
+
+    def check_keys(self, values: Dict[str, torch.Tensor], known_keys: Set[str]):
+        """Log a warning if values contains unhandled keys."""
+        # {"alpha", "bias_indices", "bias_values", "bias_size"} are hard-coded, because they are handled by
+        # `LoRALayerBase`. Sub-classes should provide the known_keys that they handled.
+        all_known_keys = known_keys | {"alpha", "bias_indices", "bias_values", "bias_size"}
+        unknown_keys = set(values.keys()) - all_known_keys
+        if unknown_keys:
+            logger.warning(
+                f"Unexpected keys found in LoRA/LyCORIS layer, model might work incorrectly! Keys: {unknown_keys}"
+            )

invokeai/backend/lora/layers/norm_layer.py

@@ -0,0 +1,36 @@
+from typing import Dict, Optional
+
+import torch
+
+from invokeai.backend.lora.layers.lora_layer_base import LoRALayerBase
+
+
+class NormLayer(LoRALayerBase):
+    # bias handled in LoRALayerBase(calc_size, to)
+    # weight: torch.Tensor
+    # bias: Optional[torch.Tensor]
+
+    def __init__(
+        self,
+        values: Dict[str, torch.Tensor],
+    ):
+        super().__init__(values)
+
+        self.weight = values["w_norm"]
+        self.bias = values.get("b_norm", None)
+
+        self.rank = None  # unscaled
+        self.check_keys(values, {"w_norm", "b_norm"})
+
+    def get_weight(self, orig_weight: torch.Tensor) -> torch.Tensor:
+        return self.weight
+
+    def calc_size(self) -> int:
+        model_size = super().calc_size()
+        model_size += self.weight.nelement() * self.weight.element_size()
+        return model_size
+
+    def to(self, device: Optional[torch.device] = None, dtype: Optional[torch.dtype] = None) -> None:
+        super().to(device=device, dtype=dtype)
+
+        self.weight = self.weight.to(device=device, dtype=dtype)

invokeai/backend/lora/layers/utils.py

@@ -0,0 +1,33 @@
+from typing import Dict
+
+import torch
+
+from invokeai.backend.lora.layers.any_lora_layer import AnyLoRALayer
+from invokeai.backend.lora.layers.full_layer import FullLayer
+from invokeai.backend.lora.layers.ia3_layer import IA3Layer
+from invokeai.backend.lora.layers.loha_layer import LoHALayer
+from invokeai.backend.lora.layers.lokr_layer import LoKRLayer
+from invokeai.backend.lora.layers.lora_layer import LoRALayer
+from invokeai.backend.lora.layers.norm_layer import NormLayer
+
+
+def any_lora_layer_from_state_dict(state_dict: Dict[str, torch.Tensor]) -> AnyLoRALayer:
+    # Detect layers according to LyCORIS detection logic(`weight_list_det`)
+    # https://github.com/KohakuBlueleaf/LyCORIS/tree/8ad8000efb79e2b879054da8c9356e6143591bad/lycoris/modules
+
+    if "lora_up.weight" in state_dict:
+        # LoRA a.k.a LoCon
+        return LoRALayer(state_dict)
+    elif "hada_w1_a" in state_dict:
+        return LoHALayer(state_dict)
+    elif "lokr_w1" in state_dict or "lokr_w1_a" in state_dict:
+        return LoKRLayer(state_dict)
+    elif "diff" in state_dict:
+        # Full a.k.a Diff
+        return FullLayer(state_dict)
+    elif "on_input" in state_dict:
+        return IA3Layer(state_dict)
+    elif "w_norm" in state_dict:
+        return NormLayer(state_dict)
+    else:
+        raise ValueError(f"Unsupported lora format: {state_dict.keys()}")

invokeai/backend/lora/lora_model_raw.py

@@ -0,0 +1,22 @@
+# Copyright (c) 2024 The InvokeAI Development team
+from typing import Dict, Optional
+
+import torch
+
+from invokeai.backend.lora.layers.any_lora_layer import AnyLoRALayer
+from invokeai.backend.raw_model import RawModel
+
+
+class LoRAModelRaw(RawModel):  # (torch.nn.Module):
+    def __init__(self, layers: Dict[str, AnyLoRALayer]):
+        self.layers = layers
+
+    def to(self, device: Optional[torch.device] = None, dtype: Optional[torch.dtype] = None) -> None:
+        for _key, layer in self.layers.items():
+            layer.to(device=device, dtype=dtype)
+
+    def calc_size(self) -> int:
+        model_size = 0
+        for _, layer in self.layers.items():
+            model_size += layer.calc_size()
+        return model_size

invokeai/backend/lora/lora_patcher.py

@@ -0,0 +1,148 @@
+from contextlib import contextmanager
+from typing import Dict, Iterable, Optional, Tuple
+
+import torch
+
+from invokeai.backend.lora.lora_model_raw import LoRAModelRaw
+from invokeai.backend.util.devices import TorchDevice
+from invokeai.backend.util.original_weights_storage import OriginalWeightsStorage
+
+
+class LoraPatcher:
+    @staticmethod
+    @torch.no_grad()
+    @contextmanager
+    def apply_lora_patches(
+        model: torch.nn.Module,
+        patches: Iterable[Tuple[LoRAModelRaw, float]],
+        prefix: str,
+        cached_weights: Optional[Dict[str, torch.Tensor]] = None,
+    ):
+        """Apply one or more LoRA patches to a model within a context manager.
+
+        :param model: The model to patch.
+        :param loras: An iterator that returns tuples of LoRA patches and associated weights. An iterator is used so
+            that the LoRA patches do not need to be loaded into memory all at once.
+        :param prefix: The keys in the patches will be filtered to only include weights with this prefix.
+        :cached_weights: Read-only copy of the model's state dict in CPU, for efficient unpatching purposes.
+        """
+        original_weights = OriginalWeightsStorage(cached_weights)
+        try:
+            for patch, patch_weight in patches:
+                LoraPatcher.apply_lora_patch(
+                    model=model,
+                    prefix=prefix,
+                    patch=patch,
+                    patch_weight=patch_weight,
+                    original_weights=original_weights,
+                )
+                del patch
+
+            yield
+        finally:
+            for param_key, weight in original_weights.get_changed_weights():
+                model.get_parameter(param_key).copy_(weight)
+
+    @staticmethod
+    @torch.no_grad()
+    def apply_lora_patch(
+        model: torch.nn.Module,
+        prefix: str,
+        patch: LoRAModelRaw,
+        patch_weight: float,
+        original_weights: OriginalWeightsStorage,
+    ):
+        """
+        Apply a single LoRA patch to a model.
+        :param model: The model to patch.
+        :param patch: LoRA model to patch in.
+        :param patch_weight: LoRA patch weight.
+        :param prefix: A string prefix that precedes keys used in the LoRAs weight layers.
+        :param original_weights: Storage with original weights, filled by weights which lora patches, used for unpatching.
+        """
+
+        if patch_weight == 0:
+            return
+
+        # If the layer keys contain a dot, then they are not flattened, and can be directly used to access model
+        # submodules. If the layer keys do not contain a dot, then they are flattened, meaning that all '.' have been
+        # replaced with '_'. Non-flattened keys are preferred, because they allow submodules to be accessed directly
+        # without searching, but some legacy code still uses flattened keys.
+        layer_keys_are_flattened = "." not in next(iter(patch.layers.keys()))
+
+        prefix_len = len(prefix)
+
+        for layer_key, layer in patch.layers.items():
+            if not layer_key.startswith(prefix):
+                continue
+
+            module_key, module = LoraPatcher._get_submodule(
+                model, layer_key[prefix_len:], layer_key_is_flattened=layer_keys_are_flattened
+            )
+
+            # All of the LoRA weight calculations will be done on the same device as the module weight.
+            # (Performance will be best if this is a CUDA device.)
+            device = module.weight.device
+            dtype = module.weight.dtype
+
+            layer_scale = layer.alpha / layer.rank if (layer.alpha and layer.rank) else 1.0
+
+            # We intentionally move to the target device first, then cast. Experimentally, this was found to
+            # be significantly faster for 16-bit CPU tensors being moved to a CUDA device than doing the
+            # same thing in a single call to '.to(...)'.
+            layer.to(device=device)
+            layer.to(dtype=torch.float32)
+
+            # TODO(ryand): Using torch.autocast(...) over explicit casting may offer a speed benefit on CUDA
+            # devices here. Experimentally, it was found to be very slow on CPU. More investigation needed.
+            for param_name, lora_param_weight in layer.get_parameters(module).items():
+                param_key = module_key + "." + param_name
+                module_param = module.get_parameter(param_name)
+
+                # Save original weight
+                original_weights.save(param_key, module_param)
+
+                if module_param.shape != lora_param_weight.shape:
+                    lora_param_weight = lora_param_weight.reshape(module_param.shape)
+
+                lora_param_weight *= patch_weight * layer_scale
+                module_param += lora_param_weight.to(dtype=dtype)
+
+            layer.to(device=TorchDevice.CPU_DEVICE)
+
+    @staticmethod
+    def _get_submodule(
+        model: torch.nn.Module, layer_key: str, layer_key_is_flattened: bool
+    ) -> tuple[str, torch.nn.Module]:
+        """Get the submodule corresponding to the given layer key.
+        :param model: The model to search.
+        :param layer_key: The layer key to search for.
+        :param layer_key_is_flattened: Whether the layer key is flattened. If flattened, then all '.' have been replaced
+            with '_'. Non-flattened keys are preferred, because they allow submodules to be accessed directly without
+            searching, but some legacy code still uses flattened keys.
+        :return: A tuple containing the module key and the submodule.
+        """
+        if not layer_key_is_flattened:
+            return layer_key, model.get_submodule(layer_key)
+
+        # Handle flattened keys.
+        assert "." not in layer_key
+
+        module = model
+        module_key = ""
+        key_parts = layer_key.split("_")
+
+        submodule_name = key_parts.pop(0)
+
+        while len(key_parts) > 0:
+            try:
+                module = module.get_submodule(submodule_name)
+                module_key += "." + submodule_name
+                submodule_name = key_parts.pop(0)
+            except Exception:
+                submodule_name += "_" + key_parts.pop(0)
+
+        module = module.get_submodule(submodule_name)
+        module_key = (module_key + "." + submodule_name).lstrip(".")
+
+        return module_key, module

invokeai/backend/model_manager/load/model_loaders/lora.py

@@ -5,8 +5,18 @@ from logging import Logger
 from pathlib import Path
 from typing import Optional
 
+import torch
+from safetensors.torch import load_file
+
 from invokeai.app.services.config import InvokeAIAppConfig
-from invokeai.backend.lora import LoRAModelRaw
+from invokeai.backend.lora.conversions.flux_diffusers_lora_conversion_utils import (
+    lora_model_from_flux_diffusers_state_dict,
+)
+from invokeai.backend.lora.conversions.flux_kohya_lora_conversion_utils import (
+    lora_model_from_flux_kohya_state_dict,
+)
+from invokeai.backend.lora.conversions.sd_lora_conversion_utils import lora_model_from_sd_state_dict
+from invokeai.backend.lora.conversions.sdxl_lora_conversion_utils import convert_sdxl_keys_to_diffusers_format
 from invokeai.backend.model_manager import (
     AnyModel,
     AnyModelConfig,
@@ -45,14 +55,33 @@ class LoRALoader(ModelLoader):
             raise ValueError("There are no submodels in a LoRA model.")
         model_path = Path(config.path)
         assert self._model_base is not None
-        model = LoRAModelRaw.from_checkpoint(
-            file_path=model_path,
-            dtype=self._torch_dtype,
-            base_model=self._model_base,
-        )
+
+        # Load the state dict from the model file.
+        if model_path.suffix == ".safetensors":
+            state_dict = load_file(model_path.absolute().as_posix(), device="cpu")
+        else:
+            state_dict = torch.load(model_path, map_location="cpu")
+
+        # Apply state_dict key conversions, if necessary.
+        if self._model_base == BaseModelType.StableDiffusionXL:
+            state_dict = convert_sdxl_keys_to_diffusers_format(state_dict)
+            model = lora_model_from_sd_state_dict(state_dict=state_dict)
+        elif self._model_base == BaseModelType.Flux:
+            if config.format == ModelFormat.Diffusers:
+                model = lora_model_from_flux_diffusers_state_dict(state_dict=state_dict)
+            elif config.format == ModelFormat.LyCORIS:
+                model = lora_model_from_flux_kohya_state_dict(state_dict=state_dict)
+            else:
+                raise ValueError(f"LoRA model is in unsupported FLUX format: {config.format}")
+        elif self._model_base in [BaseModelType.StableDiffusion1, BaseModelType.StableDiffusion2]:
+            # Currently, we don't apply any conversions for SD1 and SD2 LoRA models.
+            model = lora_model_from_sd_state_dict(state_dict=state_dict)
+        else:
+            raise ValueError(f"Unsupported LoRA base model: {self._model_base}")
+
+        model.to(dtype=self._torch_dtype)
         return model
 
-    # override
     def _get_model_path(self, config: AnyModelConfig) -> Path:
         # cheating a little - we remember this variable for using in the subsequent call to _load_model()
         self._model_base = config.base

invokeai/backend/model_manager/load/model_util.py

@@ -15,7 +15,7 @@ from invokeai.backend.image_util.depth_anything.depth_anything_pipeline import D
 from invokeai.backend.image_util.grounding_dino.grounding_dino_pipeline import GroundingDinoPipeline
 from invokeai.backend.image_util.segment_anything.segment_anything_pipeline import SegmentAnythingPipeline
 from invokeai.backend.ip_adapter.ip_adapter import IPAdapter
-from invokeai.backend.lora import LoRAModelRaw
+from invokeai.backend.lora.lora_model_raw import LoRAModelRaw
 from invokeai.backend.model_manager.config import AnyModel
 from invokeai.backend.onnx.onnx_runtime import IAIOnnxRuntimeModel
 from invokeai.backend.spandrel_image_to_image_model import SpandrelImageToImageModel

invokeai/backend/model_manager/probe.py

@@ -10,6 +10,10 @@ from picklescan.scanner import scan_file_path
 
 import invokeai.backend.util.logging as logger
 from invokeai.app.util.misc import uuid_string
+from invokeai.backend.lora.conversions.flux_diffusers_lora_conversion_utils import (
+    is_state_dict_likely_in_flux_diffusers_format,
+)
+from invokeai.backend.lora.conversions.flux_kohya_lora_conversion_utils import is_state_dict_likely_in_flux_kohya_format
 from invokeai.backend.model_hash.model_hash import HASHING_ALGORITHMS, ModelHash
 from invokeai.backend.model_manager.config import (
     AnyModelConfig,
@@ -244,7 +248,9 @@ class ModelProbe(object):
                 return ModelType.VAE
             elif key.startswith(("lora_te_", "lora_unet_")):
                 return ModelType.LoRA
-            elif key.endswith(("to_k_lora.up.weight", "to_q_lora.down.weight")):
+            # "lora_A.weight" and "lora_B.weight" are associated with models in PEFT format. We don't support all PEFT
+            # LoRA models, but as of the time of writing, we support Diffusers FLUX PEFT LoRA models.
+            elif key.endswith(("to_k_lora.up.weight", "to_q_lora.down.weight", "lora_A.weight", "lora_B.weight")):
                 return ModelType.LoRA
             elif key.startswith(("controlnet", "control_model", "input_blocks")):
                 return ModelType.ControlNet
@@ -554,12 +560,21 @@ class LoRACheckpointProbe(CheckpointProbeBase):
     """Class for LoRA checkpoints."""
 
     def get_format(self) -> ModelFormat:
-        return ModelFormat("lycoris")
+        if is_state_dict_likely_in_flux_diffusers_format(self.checkpoint):
+            # TODO(ryand): This is an unusual case. In other places throughout the codebase, we treat
+            # ModelFormat.Diffusers as meaning that the model is in a directory. In this case, the model is a single
+            # file, but the weight keys are in the diffusers format.
+            return ModelFormat.Diffusers
+        return ModelFormat.LyCORIS
 
     def get_base_type(self) -> BaseModelType:
-        checkpoint = self.checkpoint
-        token_vector_length = lora_token_vector_length(checkpoint)
+        if is_state_dict_likely_in_flux_kohya_format(self.checkpoint) or is_state_dict_likely_in_flux_diffusers_format(
+            self.checkpoint
+        ):
+            return BaseModelType.Flux
 
+        # If we've gotten here, we assume that the model is a Stable Diffusion model.
+        token_vector_length = lora_token_vector_length(self.checkpoint)
         if token_vector_length == 768:
             return BaseModelType.StableDiffusion1
         elif token_vector_length == 1024:

invokeai/backend/model_patcher.py

@@ -5,32 +5,18 @@ from __future__ import annotations
 
 import pickle
 from contextlib import contextmanager
-from typing import Any, Dict, Generator, Iterator, List, Optional, Tuple, Type, Union
+from typing import Any, Dict, Iterator, List, Optional, Tuple, Type, Union
 
 import numpy as np
 import torch
-from diffusers import OnnxRuntimeModel, UNet2DConditionModel
+from diffusers import UNet2DConditionModel
 from transformers import CLIPTextModel, CLIPTextModelWithProjection, CLIPTokenizer
 
 from invokeai.app.shared.models import FreeUConfig
-from invokeai.backend.lora import LoRAModelRaw
-from invokeai.backend.model_manager import AnyModel
+from invokeai.backend.lora.lora_model_raw import LoRAModelRaw
 from invokeai.backend.model_manager.load.optimizations import skip_torch_weight_init
 from invokeai.backend.onnx.onnx_runtime import IAIOnnxRuntimeModel
-from invokeai.backend.stable_diffusion.extensions.lora import LoRAExt
 from invokeai.backend.textual_inversion import TextualInversionManager, TextualInversionModelRaw
-from invokeai.backend.util.original_weights_storage import OriginalWeightsStorage
-
-"""
-loras = [
-    (lora_model1, 0.7),
-    (lora_model2, 0.4),
-]
-with LoRAHelper.apply_lora_unet(unet, loras):
-    # unet with applied loras
-# unmodified unet
-
-"""
 
 
 class ModelPatcher:
@@ -54,95 +40,6 @@ class ModelPatcher:
         finally:
             unet.set_attn_processor(unet_orig_processors)
 
-    @staticmethod
-    def _resolve_lora_key(model: torch.nn.Module, lora_key: str, prefix: str) -> Tuple[str, torch.nn.Module]:
-        assert "." not in lora_key
-
-        if not lora_key.startswith(prefix):
-            raise Exception(f"lora_key with invalid prefix: {lora_key}, {prefix}")
-
-        module = model
-        module_key = ""
-        key_parts = lora_key[len(prefix) :].split("_")
-
-        submodule_name = key_parts.pop(0)
-
-        while len(key_parts) > 0:
-            try:
-                module = module.get_submodule(submodule_name)
-                module_key += "." + submodule_name
-                submodule_name = key_parts.pop(0)
-            except Exception:
-                submodule_name += "_" + key_parts.pop(0)
-
-        module = module.get_submodule(submodule_name)
-        module_key = (module_key + "." + submodule_name).lstrip(".")
-
-        return (module_key, module)
-
-    @classmethod
-    @contextmanager
-    def apply_lora_unet(
-        cls,
-        unet: UNet2DConditionModel,
-        loras: Iterator[Tuple[LoRAModelRaw, float]],
-        cached_weights: Optional[Dict[str, torch.Tensor]] = None,
-    ) -> Generator[None, None, None]:
-        with cls.apply_lora(
-            unet,
-            loras=loras,
-            prefix="lora_unet_",
-            cached_weights=cached_weights,
-        ):
-            yield
-
-    @classmethod
-    @contextmanager
-    def apply_lora_text_encoder(
-        cls,
-        text_encoder: CLIPTextModel,
-        loras: Iterator[Tuple[LoRAModelRaw, float]],
-        cached_weights: Optional[Dict[str, torch.Tensor]] = None,
-    ) -> Generator[None, None, None]:
-        with cls.apply_lora(text_encoder, loras=loras, prefix="lora_te_", cached_weights=cached_weights):
-            yield
-
-    @classmethod
-    @contextmanager
-    def apply_lora(
-        cls,
-        model: AnyModel,
-        loras: Iterator[Tuple[LoRAModelRaw, float]],
-        prefix: str,
-        cached_weights: Optional[Dict[str, torch.Tensor]] = None,
-    ) -> Generator[None, None, None]:
-        """
-        Apply one or more LoRAs to a model.
-
-        :param model: The model to patch.
-        :param loras: An iterator that returns the LoRA to patch in and its patch weight.
-        :param prefix: A string prefix that precedes keys used in the LoRAs weight layers.
-        :cached_weights: Read-only copy of the model's state dict in CPU, for unpatching purposes.
-        """
-        original_weights = OriginalWeightsStorage(cached_weights)
-        try:
-            for lora_model, lora_weight in loras:
-                LoRAExt.patch_model(
-                    model=model,
-                    prefix=prefix,
-                    lora=lora_model,
-                    lora_weight=lora_weight,
-                    original_weights=original_weights,
-                )
-                del lora_model
-
-            yield
-
-        finally:
-            with torch.no_grad():
-                for param_key, weight in original_weights.get_changed_weights():
-                    model.get_parameter(param_key).copy_(weight)
-
     @classmethod
     @contextmanager
     def apply_ti(
@@ -282,26 +179,6 @@ class ModelPatcher:
 
 
 class ONNXModelPatcher:
-    @classmethod
-    @contextmanager
-    def apply_lora_unet(
-        cls,
-        unet: OnnxRuntimeModel,
-        loras: Iterator[Tuple[LoRAModelRaw, float]],
-    ) -> None:
-        with cls.apply_lora(unet, loras, "lora_unet_"):
-            yield
-
-    @classmethod
-    @contextmanager
-    def apply_lora_text_encoder(
-        cls,
-        text_encoder: OnnxRuntimeModel,
-        loras: List[Tuple[LoRAModelRaw, float]],
-    ) -> None:
-        with cls.apply_lora(text_encoder, loras, "lora_te_"):
-            yield
-
     # based on
     # https://github.com/ssube/onnx-web/blob/ca2e436f0623e18b4cfe8a0363fcfcf10508acf7/api/onnx_web/convert/diffusion/lora.py#L323
     @classmethod

invokeai/backend/stable_diffusion/extensions/lora.py

@@ -1,18 +1,17 @@
 from __future__ import annotations
 
 from contextlib import contextmanager
-from typing import TYPE_CHECKING, Tuple
+from typing import TYPE_CHECKING
 
-import torch
 from diffusers import UNet2DConditionModel
 
+from invokeai.backend.lora.lora_model_raw import LoRAModelRaw
+from invokeai.backend.lora.lora_patcher import LoraPatcher
 from invokeai.backend.stable_diffusion.extensions.base import ExtensionBase
-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.backend.lora import LoRAModelRaw
     from invokeai.backend.util.original_weights_storage import OriginalWeightsStorage
 
 
@@ -31,107 +30,14 @@ class LoRAExt(ExtensionBase):
     @contextmanager
     def patch_unet(self, unet: UNet2DConditionModel, original_weights: OriginalWeightsStorage):
         lora_model = self._node_context.models.load(self._model_id).model
-        self.patch_model(
+        assert isinstance(lora_model, LoRAModelRaw)
+        LoraPatcher.apply_lora_patch(
             model=unet,
             prefix="lora_unet_",
-            lora=lora_model,
-            lora_weight=self._weight,
+            patch=lora_model,
+            patch_weight=self._weight,
             original_weights=original_weights,
         )
         del lora_model
 
         yield
-
-    @classmethod
-    @torch.no_grad()
-    def patch_model(
-        cls,
-        model: torch.nn.Module,
-        prefix: str,
-        lora: LoRAModelRaw,
-        lora_weight: float,
-        original_weights: OriginalWeightsStorage,
-    ):
-        """
-        Apply one or more LoRAs to a model.
-        :param model: The model to patch.
-        :param lora: LoRA model to patch in.
-        :param lora_weight: LoRA patch weight.
-        :param prefix: A string prefix that precedes keys used in the LoRAs weight layers.
-        :param original_weights: Storage with original weights, filled by weights which lora patches, used for unpatching.
-        """
-
-        if lora_weight == 0:
-            return
-
-        # assert lora.device.type == "cpu"
-        for layer_key, layer in lora.layers.items():
-            if not layer_key.startswith(prefix):
-                continue
-
-            # TODO(ryand): A non-negligible amount of time is currently spent resolving LoRA keys. This
-            # should be improved in the following ways:
-            # 1. The key mapping could be more-efficiently pre-computed. This would save time every time a
-            #    LoRA model is applied.
-            # 2. From an API perspective, there's no reason that the `ModelPatcher` should be aware of the
-            #    intricacies of Stable Diffusion key resolution. It should just expect the input LoRA
-            #    weights to have valid keys.
-            assert isinstance(model, torch.nn.Module)
-            module_key, module = cls._resolve_lora_key(model, layer_key, prefix)
-
-            # All of the LoRA weight calculations will be done on the same device as the module weight.
-            # (Performance will be best if this is a CUDA device.)
-            device = module.weight.device
-            dtype = module.weight.dtype
-
-            layer_scale = layer.alpha / layer.rank if (layer.alpha and layer.rank) else 1.0
-
-            # We intentionally move to the target device first, then cast. Experimentally, this was found to
-            # be significantly faster for 16-bit CPU tensors being moved to a CUDA device than doing the
-            # same thing in a single call to '.to(...)'.
-            layer.to(device=device)
-            layer.to(dtype=torch.float32)
-
-            # TODO(ryand): Using torch.autocast(...) over explicit casting may offer a speed benefit on CUDA
-            # devices here. Experimentally, it was found to be very slow on CPU. More investigation needed.
-            for param_name, lora_param_weight in layer.get_parameters(module).items():
-                param_key = module_key + "." + param_name
-                module_param = module.get_parameter(param_name)
-
-                # save original weight
-                original_weights.save(param_key, module_param)
-
-                if module_param.shape != lora_param_weight.shape:
-                    # TODO: debug on lycoris
-                    lora_param_weight = lora_param_weight.reshape(module_param.shape)
-
-                lora_param_weight *= lora_weight * layer_scale
-                module_param += lora_param_weight.to(dtype=dtype)
-
-            layer.to(device=TorchDevice.CPU_DEVICE)
-
-    @staticmethod
-    def _resolve_lora_key(model: torch.nn.Module, lora_key: str, prefix: str) -> Tuple[str, torch.nn.Module]:
-        assert "." not in lora_key
-
-        if not lora_key.startswith(prefix):
-            raise Exception(f"lora_key with invalid prefix: {lora_key}, {prefix}")
-
-        module = model
-        module_key = ""
-        key_parts = lora_key[len(prefix) :].split("_")
-
-        submodule_name = key_parts.pop(0)
-
-        while len(key_parts) > 0:
-            try:
-                module = module.get_submodule(submodule_name)
-                module_key += "." + submodule_name
-                submodule_name = key_parts.pop(0)
-            except Exception:
-                submodule_name += "_" + key_parts.pop(0)
-
-        module = module.get_submodule(submodule_name)
-        module_key = (module_key + "." + submodule_name).lstrip(".")
-
-        return (module_key, module)

invokeai/frontend/web/.eslintrc.js

@@ -12,6 +12,18 @@ module.exports = {
     'i18next/no-literal-string': 'error',
     // https://eslint.org/docs/latest/rules/no-console
     'no-console': 'error',
+    // https://eslint.org/docs/latest/rules/no-promise-executor-return
+    'no-promise-executor-return': 'error',
+    // https://eslint.org/docs/latest/rules/require-await
+    'require-await': 'error',
+    'no-restricted-properties': [
+      'error',
+      {
+        object: 'crypto',
+        property: 'randomUUID',
+        message: 'Use of crypto.randomUUID is not allowed as it is not available in all browsers.',
+      },
+    ],
   },
   overrides: [
     /**

invokeai/frontend/web/.storybook/ReduxInit.tsx

@@ -1,5 +1,5 @@
 import { PropsWithChildren, memo, useEffect } from 'react';
-import { modelChanged } from '../src/features/parameters/store/generationSlice';
+import { modelChanged } from '../src/features/controlLayers/store/paramsSlice';
 import { useAppDispatch } from '../src/app/store/storeHooks';
 import { useGlobalModifiersInit } from '@invoke-ai/ui-library';
 /**
@@ -10,7 +10,9 @@ export const ReduxInit = memo((props: PropsWithChildren) => {
   const dispatch = useAppDispatch();
   useGlobalModifiersInit();
   useEffect(() => {
-    dispatch(modelChanged({ key: 'test_model', hash: 'some_hash', name: 'some name', base: 'sd-1', type: 'main' }));
+    dispatch(
+      modelChanged({ model: { key: 'test_model', hash: 'some_hash', name: 'some name', base: 'sd-1', type: 'main' } })
+    );
   }, []);
 
   return props.children;

invokeai/frontend/web/knip.ts

@@ -9,6 +9,10 @@ const config: KnipConfig = {
     'src/services/api/schema.ts',
     'src/features/nodes/types/v1/**',
     'src/features/nodes/types/v2/**',
+    // TODO(psyche): maybe we can clean up these utils after canvas v2 release
+    'src/features/controlLayers/konva/util.ts',
+    // TODO(psyche): restore HRF functionality?
+    'src/features/hrf/**',
   ],
   ignoreBinaries: ['only-allow'],
   paths: {

invokeai/frontend/web/package.json

@@ -24,7 +24,7 @@
     "build": "pnpm run lint && vite build",
     "typegen": "node scripts/typegen.js",
     "preview": "vite preview",
-    "lint:knip": "knip",
+    "lint:knip": "knip --tags=-knipignore",
     "lint:dpdm": "dpdm --no-warning --no-tree --transform --exit-code circular:1 src/main.tsx",
     "lint:eslint": "eslint --max-warnings=0 .",
     "lint:prettier": "prettier --check .",
@@ -52,18 +52,19 @@
     }
   },
   "dependencies": {
-    "@chakra-ui/react-use-size": "^2.1.0",
     "@dagrejs/dagre": "^1.1.3",
     "@dagrejs/graphlib": "^2.2.3",
     "@dnd-kit/core": "^6.1.0",
     "@dnd-kit/sortable": "^8.0.0",
     "@dnd-kit/utilities": "^3.2.2",
     "@fontsource-variable/inter": "^5.0.20",
-    "@invoke-ai/ui-library": "^0.0.29",
+    "@invoke-ai/ui-library": "^0.0.33",
     "@nanostores/react": "^0.7.3",
     "@reduxjs/toolkit": "2.2.3",
     "@roarr/browser-log-writer": "^1.3.0",
+    "async-mutex": "^0.5.0",
     "chakra-react-select": "^4.9.1",
+    "cmdk": "^1.0.0",
     "compare-versions": "^6.1.1",
     "dateformat": "^5.0.3",
     "fracturedjsonjs": "^4.0.2",
@@ -74,6 +75,8 @@
     "jsondiffpatch": "^0.6.0",
     "konva": "^9.3.14",
     "lodash-es": "^4.17.21",
+    "lru-cache": "^11.0.0",
+    "nanoid": "^5.0.7",
     "nanostores": "^0.11.2",
     "new-github-issue-url": "^1.0.0",
     "overlayscrollbars": "^2.10.0",
@@ -88,10 +91,8 @@
     "react-hotkeys-hook": "4.5.0",
     "react-i18next": "^14.1.3",
     "react-icons": "^5.2.1",
-    "react-konva": "^18.2.10",
     "react-redux": "9.1.2",
-    "react-resizable-panels": "^2.0.23",
-    "react-select": "5.8.0",
+    "react-resizable-panels": "^2.1.2",
     "react-use": "^17.5.1",
     "react-virtuoso": "^4.9.0",
     "reactflow": "^11.11.4",
@@ -102,9 +103,9 @@
     "roarr": "^7.21.1",
     "serialize-error": "^11.0.3",
     "socket.io-client": "^4.7.5",
+    "stable-hash": "^0.0.4",
     "use-debounce": "^10.0.2",
     "use-device-pixel-ratio": "^1.1.2",
-    "use-image": "^1.1.1",
     "uuid": "^10.0.0",
     "zod": "^3.23.8",
     "zod-validation-error": "^3.3.1"
@@ -135,6 +136,7 @@
     "@vitest/coverage-v8": "^1.5.0",
     "@vitest/ui": "^1.5.0",
     "concurrently": "^8.2.2",
+    "csstype": "^3.1.3",
     "dpdm": "^3.14.0",
     "eslint": "^8.57.0",
     "eslint-plugin-i18next": "^6.0.9",

invokeai/frontend/web/public/locales/en.json

@@ -80,6 +80,7 @@
         "aboutDesc": "Using Invoke for work? Check out:",
         "aboutHeading": "Own Your Creative Power",
         "accept": "Accept",
+        "apply": "Apply",
         "add": "Add",
         "advanced": "Advanced",
         "ai": "ai",
@@ -92,6 +93,7 @@
         "copy": "Copy",
         "copyError": "$t(gallery.copy) Error",
         "on": "On",
+        "off": "Off",
         "or": "or",
         "checkpoint": "Checkpoint",
         "communityLabel": "Community",
@@ -115,6 +117,7 @@
         "githubLabel": "Github",
         "goTo": "Go to",
         "hotkeysLabel": "Hotkeys",
+        "loadingImage": "Loading Image",
         "imageFailedToLoad": "Unable to Load Image",
         "img2img": "Image To Image",
         "inpaint": "inpaint",
@@ -132,6 +135,7 @@
         "nodes": "Workflows",
         "notInstalled": "Not $t(common.installed)",
         "openInNewTab": "Open in New Tab",
+        "openInViewer": "Open in Viewer",
         "orderBy": "Order By",
         "outpaint": "outpaint",
         "outputs": "Outputs",
@@ -162,10 +166,10 @@
         "alpha": "Alpha",
         "selected": "Selected",
         "tab": "Tab",
-        "viewing": "Viewing",
-        "viewingDesc": "Review images in a large gallery view",
-        "editing": "Editing",
-        "editingDesc": "Edit on the Control Layers canvas",
+        "view": "View",
+        "viewDesc": "Review images in a large gallery view",
+        "edit": "Edit",
+        "editDesc": "Edit on the Canvas",
         "comparing": "Comparing",
         "comparingDesc": "Comparing two images",
         "enabled": "Enabled",
@@ -325,6 +329,14 @@
         "canceled": "Canceled",
         "completedIn": "Completed in",
         "batch": "Batch",
+        "origin": "Origin",
+        "destination": "Destination",
+        "upscaling": "Upscaling",
+        "canvas": "Canvas",
+        "generation": "Generation",
+        "workflows": "Workflows",
+        "other": "Other",
+        "gallery": "Gallery",
         "batchFieldValues": "Batch Field Values",
         "item": "Item",
         "session": "Session",
@@ -363,6 +375,7 @@
         "useCache": "Use Cache"
     },
     "gallery": {
+        "gallery": "Gallery",
         "alwaysShowImageSizeBadge": "Always Show Image Size Badge",
         "assets": "Assets",
         "autoAssignBoardOnClick": "Auto-Assign Board on Click",
@@ -375,11 +388,11 @@
         "deleteImage_one": "Delete Image",
         "deleteImage_other": "Delete {{count}} Images",
         "deleteImagePermanent": "Deleted images cannot be restored.",
-        "displayBoardSearch": "Display Board Search",
-        "displaySearch": "Display Search",
+        "displayBoardSearch": "Board Search",
+        "displaySearch": "Image Search",
         "download": "Download",
         "exitBoardSearch": "Exit Board Search",
-        "exitSearch": "Exit Search",
+        "exitSearch": "Exit Image Search",
         "featuresWillReset": "If you delete this image, those features will immediately be reset.",
         "galleryImageSize": "Image Size",
         "gallerySettings": "Gallery Settings",
@@ -425,7 +438,8 @@
         "compareHelp1": "Hold <Kbd>Alt</Kbd> while clicking a gallery image or using the arrow keys to change the compare image.",
         "compareHelp2": "Press <Kbd>M</Kbd> to cycle through comparison modes.",
         "compareHelp3": "Press <Kbd>C</Kbd> to swap the compared images.",
-        "compareHelp4": "Press <Kbd>Z</Kbd> or <Kbd>Esc</Kbd> to exit."
+        "compareHelp4": "Press <Kbd>Z</Kbd> or <Kbd>Esc</Kbd> to exit.",
+        "toggleMiniViewer": "Toggle Mini Viewer"
     },
     "hotkeys": {
         "searchHotkeys": "Search Hotkeys",
@@ -1004,6 +1018,8 @@
             "noModelForControlAdapter": "Control Adapter #{{number}} has no model selected.",
             "incompatibleBaseModelForControlAdapter": "Control Adapter #{{number}} model is incompatible with main model.",
             "noModelSelected": "No model selected",
+            "canvasManagerNotLoaded": "Canvas Manager not loaded",
+            "canvasBusy": "Canvas is busy",
             "noPrompts": "No prompts generated",
             "noNodesInGraph": "No nodes in graph",
             "systemDisconnected": "System disconnected",
@@ -1035,12 +1051,11 @@
         "scaledHeight": "Scaled H",
         "scaledWidth": "Scaled W",
         "scheduler": "Scheduler",
-        "seamlessXAxis": "Seamless Tiling X Axis",
-        "seamlessYAxis": "Seamless Tiling Y Axis",
+        "seamlessXAxis": "Seamless X Axis",
+        "seamlessYAxis": "Seamless Y Axis",
         "seed": "Seed",
         "imageActions": "Image Actions",
-        "sendToImg2Img": "Send to Image to Image",
-        "sendToUnifiedCanvas": "Send To Unified Canvas",
+        "sendToCanvas": "Send To Canvas",
         "sendToUpscale": "Send To Upscale",
         "showOptionsPanel": "Show Side Panel (O or T)",
         "shuffle": "Shuffle Seed",
@@ -1100,7 +1115,6 @@
         "confirmOnDelete": "Confirm On Delete",
         "developer": "Developer",
         "displayInProgress": "Display Progress Images",
-        "enableImageDebugging": "Enable Image Debugging",
         "enableInformationalPopovers": "Enable Informational Popovers",
         "informationalPopoversDisabled": "Informational Popovers Disabled",
         "informationalPopoversDisabledDesc": "Informational popovers have been disabled. Enable them in Settings.",
@@ -1182,8 +1196,8 @@
         "problemSavingMaskDesc": "Unable to export mask",
         "prunedQueue": "Pruned Queue",
         "resetInitialImage": "Reset Initial Image",
-        "sentToImageToImage": "Sent To Image To Image",
-        "sentToUnifiedCanvas": "Sent to Unified Canvas",
+        "sentToCanvas": "Sent to Canvas",
+        "sentToUpscale": "Sent to Upscale",
         "serverError": "Server Error",
         "sessionRef": "Session: {{sessionId}}",
         "setAsCanvasInitialImage": "Set as canvas initial image",
@@ -1567,7 +1581,7 @@
         "copyToClipboard": "Copy to Clipboard",
         "cursorPosition": "Cursor Position",
         "darkenOutsideSelection": "Darken Outside Selection",
-        "discardAll": "Discard All",
+        "discardAll": "Discard All & Cancel Pending Generations",
         "discardCurrent": "Discard Current",
         "downloadAsImage": "Download As Image",
         "enableMask": "Enable Mask",
@@ -1645,39 +1659,187 @@
         "storeNotInitialized": "Store is not initialized"
     },
     "controlLayers": {
-        "deleteAll": "Delete All",
+        "bookmark": "Bookmark for Quick Switch",
+        "fitBboxToLayers": "Fit Bbox To Layers",
+        "removeBookmark": "Remove Bookmark",
+        "saveCanvasToGallery": "Save Canvas To Gallery",
+        "saveBboxToGallery": "Save Bbox To Gallery",
+        "savedToGalleryOk": "Saved to Gallery",
+        "savedToGalleryError": "Error saving to gallery",
+        "mergeVisible": "Merge Visible",
+        "mergeVisibleOk": "Merged visible layers",
+        "mergeVisibleError": "Error merging visible layers",
+        "clearHistory": "Clear History",
+        "generateMode": "Generate",
+        "generateModeDesc": "Create individual images. Generated images are added directly to the gallery.",
+        "composeMode": "Compose",
+        "composeModeDesc": "Compose your work iterative. Generated images are added back to the canvas.",
+        "autoSave": "Auto-save to Gallery",
+        "resetCanvas": "Reset Canvas",
+        "resetAll": "Reset All",
+        "clearCaches": "Clear Caches",
+        "recalculateRects": "Recalculate Rects",
+        "clipToBbox": "Clip Strokes to Bbox",
+        "compositeMaskedRegions": "Composite Masked Regions",
         "addLayer": "Add Layer",
+        "duplicate": "Duplicate",
         "moveToFront": "Move to Front",
         "moveToBack": "Move to Back",
         "moveForward": "Move Forward",
         "moveBackward": "Move Backward",
         "brushSize": "Brush Size",
+        "width": "Width",
+        "zoom": "Zoom",
+        "resetView": "Reset View",
         "controlLayers": "Control Layers",
         "globalMaskOpacity": "Global Mask Opacity",
         "autoNegative": "Auto Negative",
+        "enableAutoNegative": "Enable Auto Negative",
+        "disableAutoNegative": "Disable Auto Negative",
         "deletePrompt": "Delete Prompt",
         "resetRegion": "Reset Region",
         "debugLayers": "Debug Layers",
+        "showHUD": "Show HUD",
         "rectangle": "Rectangle",
-        "maskPreviewColor": "Mask Preview Color",
+        "maskFill": "Mask Fill",
         "addPositivePrompt": "Add $t(common.positivePrompt)",
         "addNegativePrompt": "Add $t(common.negativePrompt)",
         "addIPAdapter": "Add $t(common.ipAdapter)",
-        "regionalGuidance": "Regional Guidance",
+        "addRasterLayer": "Add $t(controlLayers.rasterLayer)",
+        "addControlLayer": "Add $t(controlLayers.controlLayer)",
+        "addInpaintMask": "Add $t(controlLayers.inpaintMask)",
+        "addRegionalGuidance": "Add $t(controlLayers.regionalGuidance)",
         "regionalGuidanceLayer": "$t(controlLayers.regionalGuidance) $t(unifiedCanvas.layer)",
+        "raster": "Raster",
+        "rasterLayer": "Raster Layer",
+        "controlLayer": "Control Layer",
+        "inpaintMask": "Inpaint Mask",
+        "regionalGuidance": "Regional Guidance",
+        "ipAdapter": "IP Adapter",
+        "sendingToCanvas": "Sending to Canvas",
+        "sendingToGallery": "Sending to Gallery",
+        "sendToGallery": "Send To Gallery",
+        "sendToGalleryDesc": "Generations will be sent to the gallery.",
+        "sendToCanvas": "Send To Canvas",
+        "sendToCanvasDesc": "Generations will be staged onto the canvas.",
+        "rasterLayer_withCount_one": "$t(controlLayers.rasterLayer)",
+        "controlLayer_withCount_one": "$t(controlLayers.controlLayer)",
+        "inpaintMask_withCount_one": "$t(controlLayers.inpaintMask)",
+        "regionalGuidance_withCount_one": "$t(controlLayers.regionalGuidance)",
+        "ipAdapter_withCount_one": "$t(controlLayers.ipAdapter)",
+        "rasterLayer_withCount_other": "Raster Layers",
+        "controlLayer_withCount_other": "Control Layers",
+        "inpaintMask_withCount_other": "Inpaint Masks",
+        "regionalGuidance_withCount_other": "Regional Guidance",
+        "ipAdapter_withCount_other": "IP Adapters",
         "opacity": "Opacity",
+        "regionalGuidance_withCount_hidden": "Regional Guidance ({{count}} hidden)",
+        "controlLayers_withCount_hidden": "Control Layers ({{count}} hidden)",
+        "rasterLayers_withCount_hidden": "Raster Layers ({{count}} hidden)",
+        "globalIPAdapters_withCount_hidden": "Global IP Adapters ({{count}} hidden)",
+        "inpaintMasks_withCount_hidden": "Inpaint Masks ({{count}} hidden)",
+        "regionalGuidance_withCount_visible": "Regional Guidance ({{count}})",
+        "controlLayers_withCount_visible": "Control Layers ({{count}})",
+        "rasterLayers_withCount_visible": "Raster Layers ({{count}})",
+        "globalIPAdapters_withCount_visible": "Global IP Adapters ({{count}})",
+        "inpaintMasks_withCount_visible": "Inpaint Masks ({{count}})",
         "globalControlAdapter": "Global $t(controlnet.controlAdapter_one)",
         "globalControlAdapterLayer": "Global $t(controlnet.controlAdapter_one) $t(unifiedCanvas.layer)",
         "globalIPAdapter": "Global $t(common.ipAdapter)",
         "globalIPAdapterLayer": "Global $t(common.ipAdapter) $t(unifiedCanvas.layer)",
         "globalInitialImage": "Global Initial Image",
         "globalInitialImageLayer": "$t(controlLayers.globalInitialImage) $t(unifiedCanvas.layer)",
+        "layer": "Layer",
         "opacityFilter": "Opacity Filter",
         "clearProcessor": "Clear Processor",
         "resetProcessor": "Reset Processor to Defaults",
         "noLayersAdded": "No Layers Added",
-        "layers_one": "Layer",
-        "layers_other": "Layers"
+        "layer_one": "Layer",
+        "layer_other": "Layers",
+        "objects_zero": "empty",
+        "objects_one": "{{count}} object",
+        "objects_other": "{{count}} objects",
+        "convertToControlLayer": "Convert to Control Layer",
+        "convertToRasterLayer": "Convert to Raster Layer",
+        "transparency": "Transparency",
+        "enableTransparencyEffect": "Enable Transparency Effect",
+        "disableTransparencyEffect": "Disable Transparency Effect",
+        "hidingType": "Hiding {{type}}",
+        "showingType": "Showing {{type}}",
+        "dynamicGrid": "Dynamic Grid",
+        "logDebugInfo": "Log Debug Info",
+        "locked": "Locked",
+        "unlocked": "Unlocked",
+        "deleteSelected": "Delete Selected",
+        "deleteAll": "Delete All",
+        "flipHorizontal": "Flip Horizontal",
+        "flipVertical": "Flip Vertical",
+        "fill": {
+            "fillColor": "Fill Color",
+            "fillStyle": "Fill Style",
+            "solid": "Solid",
+            "grid": "Grid",
+            "crosshatch": "Crosshatch",
+            "vertical": "Vertical",
+            "horizontal": "Horizontal",
+            "diagonal": "Diagonal"
+        },
+        "tool": {
+            "brush": "Brush",
+            "eraser": "Eraser",
+            "rectangle": "Rectangle",
+            "bbox": "Bbox",
+            "move": "Move",
+            "view": "View",
+            "colorPicker": "Color Picker"
+        },
+        "filter": {
+            "filter": "Filter",
+            "filters": "Filters",
+            "filterType": "Filter Type",
+            "autoProcess": "Auto Process",
+            "reset": "Reset",
+            "process": "Process",
+            "apply": "Apply",
+            "cancel": "Cancel",
+            "spandrel": {
+                "label": "Image-to-Image Model",
+                "description": "Run an image-to-image model on the selected layer.",
+                "paramModel": "Model",
+                "paramAutoScale": "Auto Scale",
+                "paramAutoScaleDesc": "The selected model will be run until the target scale is reached.",
+                "paramScale": "Target Scale"
+            }
+        },
+        "transform": {
+            "transform": "Transform",
+            "fitToBbox": "Fit to Bbox",
+            "reset": "Reset",
+            "apply": "Apply",
+            "cancel": "Cancel"
+        },
+        "settings": {
+            "snapToGrid": {
+                "label": "Snap to Grid",
+                "on": "On",
+                "off": "Off"
+            }
+        },
+        "HUD": {
+            "bbox": "Bbox",
+            "scaledBbox": "Scaled Bbox",
+            "autoSave": "Auto Save",
+            "entityStatus": {
+                "selectedEntity": "Selected Entity",
+                "selectedEntityIs": "Selected Entity is",
+                "isFiltering": "is filtering",
+                "isTransforming": "is transforming",
+                "isLocked": "is locked",
+                "isHidden": "is hidden",
+                "isDisabled": "is disabled",
+                "enabled": "Enabled"
+            }
+        }
     },
     "upscaling": {
         "upscale": "Upscale",
@@ -1765,5 +1927,30 @@
             "upscaling": "Upscaling",
             "upscalingTab": "$t(ui.tabs.upscaling) $t(common.tab)"
         }
+    },
+    "system": {
+        "enableLogging": "Enable Logging",
+        "logLevel": {
+            "logLevel": "Log Level",
+            "trace": "Trace",
+            "debug": "Debug",
+            "info": "Info",
+            "warn": "Warn",
+            "error": "Error",
+            "fatal": "Fatal"
+        },
+        "logNamespaces": {
+            "logNamespaces": "Log Namespaces",
+            "gallery": "Gallery",
+            "models": "Models",
+            "config": "Config",
+            "canvas": "Canvas",
+            "generation": "Generation",
+            "workflows": "Workflows",
+            "system": "System",
+            "events": "Events",
+            "queue": "Queue",
+            "metadata": "Metadata"
+        }
     }
 }

invokeai/frontend/web/scripts/typegen.js

@@ -38,7 +38,7 @@ async function generateTypes(schema) {
   process.stdout.write(`\nOK!\r\n`);
 }
 
-async function main() {
+function main() {
   const encoding = 'utf-8';
 
   if (process.stdin.isTTY) {

invokeai/frontend/web/src/app/components/App.tsx

@@ -6,6 +6,7 @@ import { appStarted } from 'app/store/middleware/listenerMiddleware/listeners/ap
 import { useAppDispatch, useAppSelector } from 'app/store/storeHooks';
 import type { PartialAppConfig } from 'app/types/invokeai';
 import ImageUploadOverlay from 'common/components/ImageUploadOverlay';
+import { useScopeFocusWatcher } from 'common/hooks/interactionScopes';
 import { useClearStorage } from 'common/hooks/useClearStorage';
 import { useFullscreenDropzone } from 'common/hooks/useFullscreenDropzone';
 import { useGlobalHotkeys } from 'common/hooks/useGlobalHotkeys';
@@ -13,13 +14,16 @@ import ChangeBoardModal from 'features/changeBoardModal/components/ChangeBoardMo
 import DeleteImageModal from 'features/deleteImageModal/components/DeleteImageModal';
 import { DynamicPromptsModal } from 'features/dynamicPrompts/components/DynamicPromptsPreviewModal';
 import { useStarterModelsToast } from 'features/modelManagerV2/hooks/useStarterModelsToast';
+import { ClearQueueConfirmationsAlertDialog } from 'features/queue/components/ClearQueueConfirmationAlertDialog';
 import { StylePresetModal } from 'features/stylePresets/components/StylePresetForm/StylePresetModal';
 import { activeStylePresetIdChanged } from 'features/stylePresets/store/stylePresetSlice';
+import RefreshAfterResetModal from 'features/system/components/SettingsModal/RefreshAfterResetModal';
+import SettingsModal from 'features/system/components/SettingsModal/SettingsModal';
 import { configChanged } from 'features/system/store/configSlice';
-import { languageSelector } from 'features/system/store/systemSelectors';
-import InvokeTabs from 'features/ui/components/InvokeTabs';
-import type { InvokeTabName } from 'features/ui/store/tabMap';
+import { selectLanguage } from 'features/system/store/systemSelectors';
+import { AppContent } from 'features/ui/components/AppContent';
 import { setActiveTab } from 'features/ui/store/uiSlice';
+import type { TabName } from 'features/ui/store/uiTypes';
 import { useGetAndLoadLibraryWorkflow } from 'features/workflowLibrary/hooks/useGetAndLoadLibraryWorkflow';
 import { AnimatePresence } from 'framer-motion';
 import i18n from 'i18n';
@@ -41,7 +45,7 @@ interface Props {
   };
   selectedWorkflowId?: string;
   selectedStylePresetId?: string;
-  destination?: InvokeTabName | undefined;
+  destination?: TabName;
 }
 
 const App = ({
@@ -51,7 +55,7 @@ const App = ({
   selectedStylePresetId,
   destination,
 }: Props) => {
-  const language = useAppSelector(languageSelector);
+  const language = useAppSelector(selectLanguage);
   const logger = useLogger('system');
   const dispatch = useAppDispatch();
   const clearStorage = useClearStorage();
@@ -107,6 +111,7 @@ const App = ({
 
   useStarterModelsToast();
   useSyncQueueStatus();
+  useScopeFocusWatcher();
 
   return (
     <ErrorBoundary onReset={handleReset} FallbackComponent={AppErrorBoundaryFallback}>
@@ -119,7 +124,7 @@ const App = ({
         {...dropzone.getRootProps()}
       >
         <input {...dropzone.getInputProps()} />
-        <InvokeTabs />
+        <AppContent />
         <AnimatePresence>
           {dropzone.isDragActive && isHandlingUpload && (
             <ImageUploadOverlay dropzone={dropzone} setIsHandlingUpload={setIsHandlingUpload} />
@@ -130,7 +135,10 @@ const App = ({
       <ChangeBoardModal />
       <DynamicPromptsModal />
       <StylePresetModal />
+      <ClearQueueConfirmationsAlertDialog />
       <PreselectedImage selectedImage={selectedImage} />
+      <SettingsModal />
+      <RefreshAfterResetModal />
     </ErrorBoundary>
   );
 };

invokeai/frontend/web/src/app/components/AppErrorBoundaryFallback.tsx

@@ -1,5 +1,7 @@
 import { Button, Flex, Heading, Image, Link, Text } from '@invoke-ai/ui-library';
+import { createSelector } from '@reduxjs/toolkit';
 import { useAppSelector } from 'app/store/storeHooks';
+import { selectConfigSlice } from 'features/system/store/configSlice';
 import { toast } from 'features/toast/toast';
 import newGithubIssueUrl from 'new-github-issue-url';
 import InvokeLogoYellow from 'public/assets/images/invoke-symbol-ylw-lrg.svg';
@@ -13,9 +15,11 @@ type Props = {
   resetErrorBoundary: () => void;
 };
 
+const selectIsLocal = createSelector(selectConfigSlice, (config) => config.isLocal);
+
 const AppErrorBoundaryFallback = ({ error, resetErrorBoundary }: Props) => {
   const { t } = useTranslation();
-  const isLocal = useAppSelector((s) => s.config.isLocal);
+  const isLocal = useAppSelector(selectIsLocal);
 
   const handleCopy = useCallback(() => {
     const text = JSON.stringify(serializeError(error), null, 2);

invokeai/frontend/web/src/app/components/InvokeAIUI.tsx

@@ -19,7 +19,7 @@ import type { PartialAppConfig } from 'app/types/invokeai';
 import Loading from 'common/components/Loading/Loading';
 import AppDndContext from 'features/dnd/components/AppDndContext';
 import type { WorkflowCategory } from 'features/nodes/types/workflow';
-import type { InvokeTabName } from 'features/ui/store/tabMap';
+import type { TabName } from 'features/ui/store/uiTypes';
 import type { PropsWithChildren, ReactNode } from 'react';
 import React, { lazy, memo, useEffect, useMemo } from 'react';
 import { Provider } from 'react-redux';
@@ -46,7 +46,7 @@ interface Props extends PropsWithChildren {
   };
   selectedWorkflowId?: string;
   selectedStylePresetId?: string;
-  destination?: InvokeTabName;
+  destination?: TabName;
   customStarUi?: CustomStarUi;
   socketOptions?: Partial<ManagerOptions & SocketOptions>;
   isDebugging?: boolean;

invokeai/frontend/web/src/app/components/ThemeLocaleProvider.tsx

@@ -21,10 +21,16 @@ function ThemeLocaleProvider({ children }: ThemeLocaleProviderProps) {
       direction,
       shadows: {
         ..._theme.shadows,
+        selected:
+          'inset 0px 0px 0px 3px var(--invoke-colors-invokeBlue-500), inset 0px 0px 0px 4px var(--invoke-colors-invokeBlue-800)',
+        hoverSelected:
+          'inset 0px 0px 0px 3px var(--invoke-colors-invokeBlue-400), inset 0px 0px 0px 4px var(--invoke-colors-invokeBlue-800)',
+        hoverUnselected:
+          'inset 0px 0px 0px 2px var(--invoke-colors-invokeBlue-300), inset 0px 0px 0px 3px var(--invoke-colors-invokeBlue-800)',
         selectedForCompare:
-          '0px 0px 0px 1px var(--invoke-colors-base-900), 0px 0px 0px 4px var(--invoke-colors-green-400)',
+          'inset 0px 0px 0px 3px var(--invoke-colors-invokeGreen-300), inset 0px 0px 0px 4px var(--invoke-colors-invokeGreen-800)',
         hoverSelectedForCompare:
-          '0px 0px 0px 1px var(--invoke-colors-base-900), 0px 0px 0px 4px var(--invoke-colors-green-300)',
+          'inset 0px 0px 0px 3px var(--invoke-colors-invokeGreen-200), inset 0px 0px 0px 4px var(--invoke-colors-invokeGreen-800)',
       },
     });
   }, [direction]);

invokeai/frontend/web/src/app/hooks/useSocketIO.ts

@@ -2,7 +2,7 @@ import { useStore } from '@nanostores/react';
 import { $authToken } from 'app/store/nanostores/authToken';
 import { $baseUrl } from 'app/store/nanostores/baseUrl';
 import { $isDebugging } from 'app/store/nanostores/isDebugging';
-import { useAppDispatch } from 'app/store/storeHooks';
+import { useAppStore } from 'app/store/nanostores/store';
 import type { MapStore } from 'nanostores';
 import { atom, map } from 'nanostores';
 import { useEffect, useMemo } from 'react';
@@ -18,14 +18,19 @@ declare global {
   }
 }
 
+export type AppSocket = Socket<ServerToClientEvents, ClientToServerEvents>;
+
+export const $socket = atom<AppSocket | null>(null);
 export const $socketOptions = map<Partial<ManagerOptions & SocketOptions>>({});
+
 const $isSocketInitialized = atom<boolean>(false);
+export const $isConnected = atom<boolean>(false);
 
 /**
  * Initializes the socket.io connection and sets up event listeners.
  */
 export const useSocketIO = () => {
-  const dispatch = useAppDispatch();
+  const { dispatch, getState } = useAppStore();
   const baseUrl = useStore($baseUrl);
   const authToken = useStore($authToken);
   const addlSocketOptions = useStore($socketOptions);
@@ -61,8 +66,9 @@ export const useSocketIO = () => {
       return;
     }
 
-    const socket: Socket<ServerToClientEvents, ClientToServerEvents> = io(socketUrl, socketOptions);
-    setEventListeners({ dispatch, socket });
+    const socket: AppSocket = io(socketUrl, socketOptions);
+    $socket.set(socket);
+    setEventListeners({ socket, dispatch, getState, setIsConnected: $isConnected.set });
     socket.connect();
 
     if ($isDebugging.get() || import.meta.env.MODE === 'development') {
@@ -84,5 +90,5 @@ export const useSocketIO = () => {
       socket.disconnect();
       $isSocketInitialized.set(false);
     };
-  }, [dispatch, socketOptions, socketUrl]);
+  }, [dispatch, getState, socketOptions, socketUrl]);
 };

invokeai/frontend/web/src/app/logging/logger.ts

@@ -15,21 +15,21 @@ export const BASE_CONTEXT = {};
 
 export const $logger = atom<Logger>(Roarr.child(BASE_CONTEXT));
 
-export type LoggerNamespace =
-  | 'images'
-  | 'models'
-  | 'config'
-  | 'canvas'
-  | 'generation'
-  | 'nodes'
-  | 'system'
-  | 'socketio'
-  | 'session'
-  | 'queue'
-  | 'dnd'
-  | 'controlLayers';
+export const zLogNamespace = z.enum([
+  'canvas',
+  'config',
+  'events',
+  'gallery',
+  'generation',
+  'metadata',
+  'models',
+  'system',
+  'queue',
+  'workflows',
+]);
+export type LogNamespace = z.infer<typeof zLogNamespace>;
 
-export const logger = (namespace: LoggerNamespace) => $logger.get().child({ namespace });
+export const logger = (namespace: LogNamespace) => $logger.get().child({ namespace });
 
 export const zLogLevel = z.enum(['trace', 'debug', 'info', 'warn', 'error', 'fatal']);
 export type LogLevel = z.infer<typeof zLogLevel>;

invokeai/frontend/web/src/app/logging/useLogger.ts

@@ -1,29 +1,41 @@
 import { createLogWriter } from '@roarr/browser-log-writer';
 import { useAppSelector } from 'app/store/storeHooks';
+import {
+  selectSystemLogIsEnabled,
+  selectSystemLogLevel,
+  selectSystemLogNamespaces,
+} from 'features/system/store/systemSlice';
 import { useEffect, useMemo } from 'react';
 import { ROARR, Roarr } from 'roarr';
 
-import type { LoggerNamespace } from './logger';
+import type { LogNamespace } from './logger';
 import { $logger, BASE_CONTEXT, LOG_LEVEL_MAP, logger } from './logger';
 
-export const useLogger = (namespace: LoggerNamespace) => {
-  const consoleLogLevel = useAppSelector((s) => s.system.consoleLogLevel);
-  const shouldLogToConsole = useAppSelector((s) => s.system.shouldLogToConsole);
+export const useLogger = (namespace: LogNamespace) => {
+  const logLevel = useAppSelector(selectSystemLogLevel);
+  const logNamespaces = useAppSelector(selectSystemLogNamespaces);
+  const logIsEnabled = useAppSelector(selectSystemLogIsEnabled);
 
   // The provided Roarr browser log writer uses localStorage to config logging to console
   useEffect(() => {
-    if (shouldLogToConsole) {
+    if (logIsEnabled) {
       // Enable console log output
       localStorage.setItem('ROARR_LOG', 'true');
 
       // Use a filter to show only logs of the given level
-      localStorage.setItem('ROARR_FILTER', `context.logLevel:>=${LOG_LEVEL_MAP[consoleLogLevel]}`);
+      let filter = `context.logLevel:>=${LOG_LEVEL_MAP[logLevel]}`;
+      if (logNamespaces.length > 0) {
+        filter += ` AND (${logNamespaces.map((ns) => `context.namespace:${ns}`).join(' OR ')})`;
+      } else {
+        filter += ' AND context.namespace:undefined';
+      }
+      localStorage.setItem('ROARR_FILTER', filter);
     } else {
       // Disable console log output
       localStorage.setItem('ROARR_LOG', 'false');
     }
     ROARR.write = createLogWriter();
-  }, [consoleLogLevel, shouldLogToConsole]);
+  }, [logLevel, logIsEnabled, logNamespaces]);
 
   // Update the module-scoped logger context as needed
   useEffect(() => {

invokeai/frontend/web/src/app/store/actions.ts

@@ -1,7 +1,7 @@
 import { createAction } from '@reduxjs/toolkit';
-import type { InvokeTabName } from 'features/ui/store/tabMap';
+import type { TabName } from 'features/ui/store/uiTypes';
 
 export const enqueueRequested = createAction<{
-  tabName: InvokeTabName;
+  tabName: TabName;
   prepend: boolean;
 }>('app/enqueueRequested');

invokeai/frontend/web/src/app/store/constants.ts

@@ -1,2 +1,3 @@
 export const STORAGE_PREFIX = '@@invokeai-';
 export const EMPTY_ARRAY = [];
+export const EMPTY_OBJECT = {};

invokeai/frontend/web/src/app/store/createMemoizedSelector.ts

@@ -1,5 +1,6 @@
 import { createDraftSafeSelectorCreator, createSelectorCreator, lruMemoize } from '@reduxjs/toolkit';
 import type { GetSelectorsOptions } from '@reduxjs/toolkit/dist/entities/state_selectors';
+import type { RootState } from 'app/store/store';
 import { isEqual } from 'lodash-es';
 
 /**
@@ -19,3 +20,5 @@ export const getSelectorsOptions: GetSelectorsOptions = {
     argsMemoize: lruMemoize,
   }),
 };
+
+export const createMemoizedAppSelector = createMemoizedSelector.withTypes<RootState>();

invokeai/frontend/web/src/app/store/enhancers/reduxRemember/errors.ts

@@ -1,5 +1,4 @@
 import { logger } from 'app/logging/logger';
-import { parseify } from 'common/util/serialize';
 import { PersistError, RehydrateError } from 'redux-remember';
 import { serializeError } from 'serialize-error';
 
@@ -41,6 +40,6 @@ export const errorHandler = (err: PersistError | RehydrateError) => {
   } else if (err instanceof RehydrateError) {
     log.error({ error: serializeError(err) }, 'Problem rehydrating state');
   } else {
-    log.error({ error: parseify(err) }, 'Problem in persistence layer');
+    log.error({ error: serializeError(err) }, 'Problem in persistence layer');
   }
 };

invokeai/frontend/web/src/app/store/middleware/devtools/actionSanitizer.ts

@@ -1,9 +1,7 @@
 import type { UnknownAction } from '@reduxjs/toolkit';
-import { deepClone } from 'common/util/deepClone';
 import { isAnyGraphBuilt } from 'features/nodes/store/actions';
 import { appInfoApi } from 'services/api/endpoints/appInfo';
 import type { Graph } from 'services/api/types';
-import { socketGeneratorProgress } from 'services/events/actions';
 
 export const actionSanitizer = <A extends UnknownAction>(action: A): A => {
   if (isAnyGraphBuilt(action)) {
@@ -24,13 +22,5 @@ export const actionSanitizer = <A extends UnknownAction>(action: A): A => {
     };
   }
 
-  if (socketGeneratorProgress.match(action)) {
-    const sanitized = deepClone(action);
-    if (sanitized.payload.data.progress_image) {
-      sanitized.payload.data.progress_image.dataURL = '<Progress image omitted>';
-    }
-    return sanitized;
-  }
-
   return action;
 };

invokeai/frontend/web/src/app/store/middleware/listenerMiddleware/index.ts

@@ -1,7 +1,7 @@
 import type { TypedStartListening } from '@reduxjs/toolkit';
-import { createListenerMiddleware } from '@reduxjs/toolkit';
+import { addListener, createListenerMiddleware } from '@reduxjs/toolkit';
 import { addAdHocPostProcessingRequestedListener } from 'app/store/middleware/listenerMiddleware/listeners/addAdHocPostProcessingRequestedListener';
-import { addCommitStagingAreaImageListener } from 'app/store/middleware/listenerMiddleware/listeners/addCommitStagingAreaImageListener';
+import { addStagingListeners } from 'app/store/middleware/listenerMiddleware/listeners/addCommitStagingAreaImageListener';
 import { addAnyEnqueuedListener } from 'app/store/middleware/listenerMiddleware/listeners/anyEnqueued';
 import { addAppConfigReceivedListener } from 'app/store/middleware/listenerMiddleware/listeners/appConfigReceived';
 import { addAppStartedListener } from 'app/store/middleware/listenerMiddleware/listeners/appStarted';
@@ -9,17 +9,7 @@ import { addBatchEnqueuedListener } from 'app/store/middleware/listenerMiddlewar
 import { addDeleteBoardAndImagesFulfilledListener } from 'app/store/middleware/listenerMiddleware/listeners/boardAndImagesDeleted';
 import { addBoardIdSelectedListener } from 'app/store/middleware/listenerMiddleware/listeners/boardIdSelected';
 import { addBulkDownloadListeners } from 'app/store/middleware/listenerMiddleware/listeners/bulkDownload';
-import { addCanvasCopiedToClipboardListener } from 'app/store/middleware/listenerMiddleware/listeners/canvasCopiedToClipboard';
-import { addCanvasDownloadedAsImageListener } from 'app/store/middleware/listenerMiddleware/listeners/canvasDownloadedAsImage';
-import { addCanvasImageToControlNetListener } from 'app/store/middleware/listenerMiddleware/listeners/canvasImageToControlNet';
-import { addCanvasMaskSavedToGalleryListener } from 'app/store/middleware/listenerMiddleware/listeners/canvasMaskSavedToGallery';
-import { addCanvasMaskToControlNetListener } from 'app/store/middleware/listenerMiddleware/listeners/canvasMaskToControlNet';
-import { addCanvasMergedListener } from 'app/store/middleware/listenerMiddleware/listeners/canvasMerged';
-import { addCanvasSavedToGalleryListener } from 'app/store/middleware/listenerMiddleware/listeners/canvasSavedToGallery';
-import { addControlAdapterPreprocessor } from 'app/store/middleware/listenerMiddleware/listeners/controlAdapterPreprocessor';
-import { addControlNetAutoProcessListener } from 'app/store/middleware/listenerMiddleware/listeners/controlNetAutoProcess';
-import { addControlNetImageProcessedListener } from 'app/store/middleware/listenerMiddleware/listeners/controlNetImageProcessed';
-import { addEnqueueRequestedCanvasListener } from 'app/store/middleware/listenerMiddleware/listeners/enqueueRequestedCanvas';
+import { addCancellationsListeners } from 'app/store/middleware/listenerMiddleware/listeners/cancellationsListeners';
 import { addEnqueueRequestedLinear } from 'app/store/middleware/listenerMiddleware/listeners/enqueueRequestedLinear';
 import { addEnqueueRequestedNodes } from 'app/store/middleware/listenerMiddleware/listeners/enqueueRequestedNodes';
 import { addGalleryImageClickedListener } from 'app/store/middleware/listenerMiddleware/listeners/galleryImageClicked';
@@ -37,16 +27,7 @@ import { addModelSelectedListener } from 'app/store/middleware/listenerMiddlewar
 import { addModelsLoadedListener } from 'app/store/middleware/listenerMiddleware/listeners/modelsLoaded';
 import { addDynamicPromptsListener } from 'app/store/middleware/listenerMiddleware/listeners/promptChanged';
 import { addSetDefaultSettingsListener } from 'app/store/middleware/listenerMiddleware/listeners/setDefaultSettings';
-import { addSocketConnectedEventListener } from 'app/store/middleware/listenerMiddleware/listeners/socketio/socketConnected';
-import { addSocketDisconnectedEventListener } from 'app/store/middleware/listenerMiddleware/listeners/socketio/socketDisconnected';
-import { addGeneratorProgressEventListener } from 'app/store/middleware/listenerMiddleware/listeners/socketio/socketGeneratorProgress';
-import { addInvocationCompleteEventListener } from 'app/store/middleware/listenerMiddleware/listeners/socketio/socketInvocationComplete';
-import { addInvocationErrorEventListener } from 'app/store/middleware/listenerMiddleware/listeners/socketio/socketInvocationError';
-import { addInvocationStartedEventListener } from 'app/store/middleware/listenerMiddleware/listeners/socketio/socketInvocationStarted';
-import { addModelInstallEventListener } from 'app/store/middleware/listenerMiddleware/listeners/socketio/socketModelInstall';
-import { addModelLoadEventListener } from 'app/store/middleware/listenerMiddleware/listeners/socketio/socketModelLoad';
-import { addSocketQueueItemStatusChangedEventListener } from 'app/store/middleware/listenerMiddleware/listeners/socketio/socketQueueItemStatusChanged';
-import { addStagingAreaImageSavedListener } from 'app/store/middleware/listenerMiddleware/listeners/stagingAreaImageSaved';
+import { addSocketConnectedEventListener } from 'app/store/middleware/listenerMiddleware/listeners/socketConnected';
 import { addUpdateAllNodesRequestedListener } from 'app/store/middleware/listenerMiddleware/listeners/updateAllNodesRequested';
 import { addWorkflowLoadRequestedListener } from 'app/store/middleware/listenerMiddleware/listeners/workflowLoadRequested';
 import type { AppDispatch, RootState } from 'app/store/store';
@@ -60,6 +41,8 @@ export type AppStartListening = TypedStartListening<RootState, AppDispatch>;
 
 const startAppListening = listenerMiddleware.startListening as AppStartListening;
 
+export const addAppListener = addListener.withTypes<RootState, AppDispatch>();
+
 /**
  * The RTK listener middleware is a lightweight alternative sagas/observables.
  *
@@ -83,7 +66,6 @@ addGalleryImageClickedListener(startAppListening);
 addGalleryOffsetChangedListener(startAppListening);
 
 // User Invoked
-addEnqueueRequestedCanvasListener(startAppListening);
 addEnqueueRequestedNodes(startAppListening);
 addEnqueueRequestedLinear(startAppListening);
 addEnqueueRequestedUpscale(startAppListening);
@@ -91,31 +73,22 @@ addAnyEnqueuedListener(startAppListening);
 addBatchEnqueuedListener(startAppListening);
 
 // Canvas actions
-addCanvasSavedToGalleryListener(startAppListening);
-addCanvasMaskSavedToGalleryListener(startAppListening);
-addCanvasImageToControlNetListener(startAppListening);
-addCanvasMaskToControlNetListener(startAppListening);
-addCanvasDownloadedAsImageListener(startAppListening);
-addCanvasCopiedToClipboardListener(startAppListening);
-addCanvasMergedListener(startAppListening);
-addStagingAreaImageSavedListener(startAppListening);
-addCommitStagingAreaImageListener(startAppListening);
+// addCanvasSavedToGalleryListener(startAppListening);
+// addCanvasMaskSavedToGalleryListener(startAppListening);
+// addCanvasImageToControlNetListener(startAppListening);
+// addCanvasMaskToControlNetListener(startAppListening);
+// addCanvasDownloadedAsImageListener(startAppListening);
+// addCanvasCopiedToClipboardListener(startAppListening);
+// addCanvasMergedListener(startAppListening);
+// addStagingAreaImageSavedListener(startAppListening);
+// addCommitStagingAreaImageListener(startAppListening);
+addStagingListeners(startAppListening);
 
 // Socket.IO
-addGeneratorProgressEventListener(startAppListening);
-addInvocationCompleteEventListener(startAppListening);
-addInvocationErrorEventListener(startAppListening);
-addInvocationStartedEventListener(startAppListening);
 addSocketConnectedEventListener(startAppListening);
-addSocketDisconnectedEventListener(startAppListening);
-addModelLoadEventListener(startAppListening);
-addModelInstallEventListener(startAppListening);
-addSocketQueueItemStatusChangedEventListener(startAppListening);
-addBulkDownloadListeners(startAppListening);
 
-// ControlNet
-addControlNetImageProcessedListener(startAppListening);
-addControlNetAutoProcessListener(startAppListening);
+// Gallery bulk download
+addBulkDownloadListeners(startAppListening);
 
 // Boards
 addImageAddedToBoardFulfilledListener(startAppListening);
@@ -148,4 +121,6 @@ addAdHocPostProcessingRequestedListener(startAppListening);
 addDynamicPromptsListener(startAppListening);
 
 addSetDefaultSettingsListener(startAppListening);
-addControlAdapterPreprocessor(startAppListening);
+// addControlAdapterPreprocessor(startAppListening);
+
+addCancellationsListeners(startAppListening);

invokeai/frontend/web/src/app/store/middleware/listenerMiddleware/listeners/addAdHocPostProcessingRequestedListener.ts

@@ -1,21 +1,21 @@
 import { createAction } from '@reduxjs/toolkit';
 import { logger } from 'app/logging/logger';
 import type { AppStartListening } from 'app/store/middleware/listenerMiddleware';
-import { parseify } from 'common/util/serialize';
+import type { SerializableObject } from 'common/types';
 import { buildAdHocPostProcessingGraph } from 'features/nodes/util/graph/buildAdHocPostProcessingGraph';
 import { toast } from 'features/toast/toast';
 import { t } from 'i18next';
 import { queueApi } from 'services/api/endpoints/queue';
 import type { BatchConfig, ImageDTO } from 'services/api/types';
 
+const log = logger('queue');
+
 export const adHocPostProcessingRequested = createAction<{ imageDTO: ImageDTO }>(`upscaling/postProcessingRequested`);
 
 export const addAdHocPostProcessingRequestedListener = (startAppListening: AppStartListening) => {
   startAppListening({
     actionCreator: adHocPostProcessingRequested,
     effect: async (action, { dispatch, getState }) => {
-      const log = logger('session');
-
       const { imageDTO } = action.payload;
       const state = getState();
 
@@ -39,9 +39,9 @@ export const addAdHocPostProcessingRequestedListener = (startAppListening: AppSt
 
         const enqueueResult = await req.unwrap();
         req.reset();
-        log.debug({ enqueueResult: parseify(enqueueResult) }, t('queue.graphQueued'));
+        log.debug({ enqueueResult } as SerializableObject, t('queue.graphQueued'));
       } catch (error) {
-        log.error({ enqueueBatchArg: parseify(enqueueBatchArg) }, t('queue.graphFailedToQueue'));
+        log.error({ enqueueBatchArg } as SerializableObject, t('queue.graphFailedToQueue'));
 
         if (error instanceof Object && 'status' in error && error.status === 403) {
           return;

invokeai/frontend/web/src/app/store/middleware/listenerMiddleware/listeners/addArchivedOrDeletedBoardListener.ts

@@ -23,7 +23,7 @@ export const addArchivedOrDeletedBoardListener = (startAppListening: AppStartLis
    */
   startAppListening({
     matcher: matchAnyBoardDeleted,
-    effect: async (action, { dispatch, getState }) => {
+    effect: (action, { dispatch, getState }) => {
       const state = getState();
       const deletedBoardId = action.meta.arg.originalArgs;
       const { autoAddBoardId, selectedBoardId } = state.gallery;
@@ -44,7 +44,7 @@ export const addArchivedOrDeletedBoardListener = (startAppListening: AppStartLis
   // If we archived a board, it may end up hidden. If it's selected or the auto-add board, we should reset those.
   startAppListening({
     matcher: boardsApi.endpoints.updateBoard.matchFulfilled,
-    effect: async (action, { dispatch, getState }) => {
+    effect: (action, { dispatch, getState }) => {
       const state = getState();
       const { shouldShowArchivedBoards } = state.gallery;
 
@@ -61,7 +61,7 @@ export const addArchivedOrDeletedBoardListener = (startAppListening: AppStartLis
   // When we hide archived boards, if the selected or the auto-add board is archived, we should reset those.
   startAppListening({
     actionCreator: shouldShowArchivedBoardsChanged,
-    effect: async (action, { dispatch, getState }) => {
+    effect: (action, { dispatch, getState }) => {
       const shouldShowArchivedBoards = action.payload;
 
       // We only need to take action if we have just hidden archived boards.
@@ -100,7 +100,7 @@ export const addArchivedOrDeletedBoardListener = (startAppListening: AppStartLis
    */
   startAppListening({
     matcher: boardsApi.endpoints.listAllBoards.matchFulfilled,
-    effect: async (action, { dispatch, getState }) => {
+    effect: (action, { dispatch, getState }) => {
       const boards = action.payload;
       const state = getState();
       const { selectedBoardId, autoAddBoardId } = state.gallery;

invokeai/frontend/web/src/app/store/middleware/listenerMiddleware/listeners/addCommitStagingAreaImageListener.ts

@@ -1,33 +1,34 @@
 import { isAnyOf } from '@reduxjs/toolkit';
 import { logger } from 'app/logging/logger';
 import type { AppStartListening } from 'app/store/middleware/listenerMiddleware';
-import {
-  canvasBatchIdsReset,
-  commitStagingAreaImage,
-  discardStagedImages,
-  resetCanvas,
-  setInitialCanvasImage,
-} from 'features/canvas/store/canvasSlice';
+import { canvasReset, rasterLayerAdded } from 'features/controlLayers/store/canvasSlice';
+import { stagingAreaImageAccepted, stagingAreaReset } from 'features/controlLayers/store/canvasStagingAreaSlice';
+import { selectCanvasSlice } from 'features/controlLayers/store/selectors';
+import type { CanvasRasterLayerState } from 'features/controlLayers/store/types';
+import { imageDTOToImageObject } from 'features/controlLayers/store/types';
 import { toast } from 'features/toast/toast';
 import { t } from 'i18next';
 import { queueApi } from 'services/api/endpoints/queue';
+import { assert } from 'tsafe';
 
-const matcher = isAnyOf(commitStagingAreaImage, discardStagedImages, resetCanvas, setInitialCanvasImage);
+const log = logger('canvas');
 
-export const addCommitStagingAreaImageListener = (startAppListening: AppStartListening) => {
+const matchCanvasOrStagingAreaRest = isAnyOf(stagingAreaReset, canvasReset);
+
+export const addStagingListeners = (startAppListening: AppStartListening) => {
   startAppListening({
-    matcher,
-    effect: async (_, { dispatch, getState }) => {
-      const log = logger('canvas');
-      const state = getState();
-      const { batchIds } = state.canvas;
-
+    matcher: matchCanvasOrStagingAreaRest,
+    effect: async (_, { dispatch }) => {
       try {
         const req = dispatch(
-          queueApi.endpoints.cancelByBatchIds.initiate({ batch_ids: batchIds }, { fixedCacheKey: 'cancelByBatchIds' })
+          queueApi.endpoints.cancelByBatchDestination.initiate(
+            { destination: 'canvas' },
+            { fixedCacheKey: 'cancelByBatchOrigin' }
+          )
         );
         const { canceled } = await req.unwrap();
         req.reset();
+
         if (canceled > 0) {
           log.debug(`Canceled ${canceled} canvas batches`);
           toast({
@@ -36,7 +37,6 @@ export const addCommitStagingAreaImageListener = (startAppListening: AppStartLis
             status: 'success',
           });
         }
-        dispatch(canvasBatchIdsReset());
       } catch {
         log.error('Failed to cancel canvas batches');
         toast({
@@ -47,4 +47,26 @@ export const addCommitStagingAreaImageListener = (startAppListening: AppStartLis
       }
     },
   });
+
+  startAppListening({
+    actionCreator: stagingAreaImageAccepted,
+    effect: (action, api) => {
+      const { index } = action.payload;
+      const state = api.getState();
+      const stagingAreaImage = state.canvasStagingArea.stagedImages[index];
+
+      assert(stagingAreaImage, 'No staged image found to accept');
+      const { x, y } = selectCanvasSlice(state).bbox.rect;
+
+      const { imageDTO, offsetX, offsetY } = stagingAreaImage;
+      const imageObject = imageDTOToImageObject(imageDTO);
+      const overrides: Partial<CanvasRasterLayerState> = {
+        position: { x: x + offsetX, y: y + offsetY },
+        objects: [imageObject],
+      };
+
+      api.dispatch(rasterLayerAdded({ overrides, isSelected: false }));
+      api.dispatch(stagingAreaReset());
+    },
+  });
 };

invokeai/frontend/web/src/app/store/middleware/listenerMiddleware/listeners/anyEnqueued.ts

@@ -4,7 +4,7 @@ import { queueApi, selectQueueStatus } from 'services/api/endpoints/queue';
 export const addAnyEnqueuedListener = (startAppListening: AppStartListening) => {
   startAppListening({
     matcher: queueApi.endpoints.enqueueBatch.matchFulfilled,
-    effect: async (_, { dispatch, getState }) => {
+    effect: (_, { dispatch, getState }) => {
       const { data } = selectQueueStatus(getState());
 
       if (!data || data.processor.is_started) {

invokeai/frontend/web/src/app/store/middleware/listenerMiddleware/listeners/appConfigReceived.ts

@@ -1,14 +1,14 @@
 import type { AppStartListening } from 'app/store/middleware/listenerMiddleware';
-import { setInfillMethod } from 'features/parameters/store/generationSlice';
+import { setInfillMethod } from 'features/controlLayers/store/paramsSlice';
 import { shouldUseNSFWCheckerChanged, shouldUseWatermarkerChanged } from 'features/system/store/systemSlice';
 import { appInfoApi } from 'services/api/endpoints/appInfo';
 
 export const addAppConfigReceivedListener = (startAppListening: AppStartListening) => {
   startAppListening({
     matcher: appInfoApi.endpoints.getAppConfig.matchFulfilled,
-    effect: async (action, { getState, dispatch }) => {
+    effect: (action, { getState, dispatch }) => {
       const { infill_methods = [], nsfw_methods = [], watermarking_methods = [] } = action.payload;
-      const infillMethod = getState().generation.infillMethod;
+      const infillMethod = getState().params.infillMethod;
 
       if (!infill_methods.includes(infillMethod)) {
         // if there is no infill method, set it to the first one

invokeai/frontend/web/src/app/store/middleware/listenerMiddleware/listeners/appStarted.ts

@@ -6,7 +6,7 @@ export const appStarted = createAction('app/appStarted');
 export const addAppStartedListener = (startAppListening: AppStartListening) => {
   startAppListening({
     actionCreator: appStarted,
-    effect: async (action, { unsubscribe, cancelActiveListeners }) => {
+    effect: (action, { unsubscribe, cancelActiveListeners }) => {
       // this should only run once
       cancelActiveListeners();
       unsubscribe();

invokeai/frontend/web/src/app/store/middleware/listenerMiddleware/listeners/batchEnqueued.ts

@@ -1,27 +1,30 @@
 import { logger } from 'app/logging/logger';
 import type { AppStartListening } from 'app/store/middleware/listenerMiddleware';
-import { parseify } from 'common/util/serialize';
+import type { SerializableObject } from 'common/types';
 import { zPydanticValidationError } from 'features/system/store/zodSchemas';
 import { toast } from 'features/toast/toast';
 import { t } from 'i18next';
 import { truncate, upperFirst } from 'lodash-es';
+import { serializeError } from 'serialize-error';
 import { queueApi } from 'services/api/endpoints/queue';
 
+const log = logger('queue');
+
 export const addBatchEnqueuedListener = (startAppListening: AppStartListening) => {
   // success
   startAppListening({
     matcher: queueApi.endpoints.enqueueBatch.matchFulfilled,
-    effect: async (action) => {
-      const response = action.payload;
+    effect: (action) => {
+      const enqueueResult = action.payload;
       const arg = action.meta.arg.originalArgs;
-      logger('queue').debug({ enqueueResult: parseify(response) }, 'Batch enqueued');
+      log.debug({ enqueueResult } as SerializableObject, 'Batch enqueued');
 
       toast({
         id: 'QUEUE_BATCH_SUCCEEDED',
         title: t('queue.batchQueued'),
         status: 'success',
         description: t('queue.batchQueuedDesc', {
-          count: response.enqueued,
+          count: enqueueResult.enqueued,
           direction: arg.prepend ? t('queue.front') : t('queue.back'),
         }),
       });
@@ -31,9 +34,9 @@ export const addBatchEnqueuedListener = (startAppListening: AppStartListening) =
   // error
   startAppListening({
     matcher: queueApi.endpoints.enqueueBatch.matchRejected,
-    effect: async (action) => {
+    effect: (action) => {
       const response = action.payload;
-      const arg = action.meta.arg.originalArgs;
+      const batchConfig = action.meta.arg.originalArgs;
 
       if (!response) {
         toast({
@@ -42,7 +45,7 @@ export const addBatchEnqueuedListener = (startAppListening: AppStartListening) =
           status: 'error',
           description: t('common.unknownError'),
         });
-        logger('queue').error({ batchConfig: parseify(arg), error: parseify(response) }, t('queue.batchFailedToQueue'));
+        log.error({ batchConfig } as SerializableObject, t('queue.batchFailedToQueue'));
         return;
       }
 
@@ -68,7 +71,7 @@ export const addBatchEnqueuedListener = (startAppListening: AppStartListening) =
           description: t('common.unknownError'),
         });
       }
-      logger('queue').error({ batchConfig: parseify(arg), error: parseify(response) }, t('queue.batchFailedToQueue'));
+      log.error({ batchConfig, error: serializeError(response) } as SerializableObject, t('queue.batchFailedToQueue'));
     },
   });
 };

invokeai/frontend/web/src/app/store/middleware/listenerMiddleware/listeners/boardAndImagesDeleted.ts

@@ -1,47 +1,31 @@
 import type { AppStartListening } from 'app/store/middleware/listenerMiddleware';
-import { resetCanvas } from 'features/canvas/store/canvasSlice';
-import { controlAdaptersReset } from 'features/controlAdapters/store/controlAdaptersSlice';
-import { allLayersDeleted } from 'features/controlLayers/store/controlLayersSlice';
+import { selectCanvasSlice } from 'features/controlLayers/store/selectors';
 import { getImageUsage } from 'features/deleteImageModal/store/selectors';
 import { nodeEditorReset } from 'features/nodes/store/nodesSlice';
+import { selectNodesSlice } from 'features/nodes/store/selectors';
 import { imagesApi } from 'services/api/endpoints/images';
 
 export const addDeleteBoardAndImagesFulfilledListener = (startAppListening: AppStartListening) => {
   startAppListening({
     matcher: imagesApi.endpoints.deleteBoardAndImages.matchFulfilled,
-    effect: async (action, { dispatch, getState }) => {
+    effect: (action, { dispatch, getState }) => {
       const { deleted_images } = action.payload;
 
       // Remove all deleted images from the UI
 
-      let wasCanvasReset = false;
       let wasNodeEditorReset = false;
-      let wereControlAdaptersReset = false;
-      let wereControlLayersReset = false;
 
-      const { canvas, nodes, controlAdapters, controlLayers } = getState();
+      const state = getState();
+      const nodes = selectNodesSlice(state);
+      const canvas = selectCanvasSlice(state);
+
       deleted_images.forEach((image_name) => {
-        const imageUsage = getImageUsage(canvas, nodes.present, controlAdapters, controlLayers.present, image_name);
-
-        if (imageUsage.isCanvasImage && !wasCanvasReset) {
-          dispatch(resetCanvas());
-          wasCanvasReset = true;
-        }
+        const imageUsage = getImageUsage(nodes, canvas, image_name);
 
         if (imageUsage.isNodesImage && !wasNodeEditorReset) {
           dispatch(nodeEditorReset());
           wasNodeEditorReset = true;
         }
-
-        if (imageUsage.isControlImage && !wereControlAdaptersReset) {
-          dispatch(controlAdaptersReset());
-          wereControlAdaptersReset = true;
-        }
-
-        if (imageUsage.isControlLayerImage && !wereControlLayersReset) {
-          dispatch(allLayersDeleted());
-          wereControlLayersReset = true;
-        }
       });
     },
   });

invokeai/frontend/web/src/app/store/middleware/listenerMiddleware/listeners/bulkDownload.tsx

@@ -1,21 +1,15 @@
-import { ExternalLink } from '@invoke-ai/ui-library';
 import { logger } from 'app/logging/logger';
 import type { AppStartListening } from 'app/store/middleware/listenerMiddleware';
 import { toast } from 'features/toast/toast';
 import { t } from 'i18next';
 import { imagesApi } from 'services/api/endpoints/images';
-import {
-  socketBulkDownloadComplete,
-  socketBulkDownloadError,
-  socketBulkDownloadStarted,
-} from 'services/events/actions';
 
-const log = logger('images');
+const log = logger('gallery');
 
 export const addBulkDownloadListeners = (startAppListening: AppStartListening) => {
   startAppListening({
     matcher: imagesApi.endpoints.bulkDownloadImages.matchFulfilled,
-    effect: async (action) => {
+    effect: (action) => {
       log.debug(action.payload, 'Bulk download requested');
 
       // If we have an item name, we are processing the bulk download locally and should use it as the toast id to
@@ -33,7 +27,7 @@ export const addBulkDownloadListeners = (startAppListening: AppStartListening) =
 
   startAppListening({
     matcher: imagesApi.endpoints.bulkDownloadImages.matchRejected,
-    effect: async () => {
+    effect: () => {
       log.debug('Bulk download request failed');
 
       // There isn't any toast to update if we get this event.
@@ -44,55 +38,4 @@ export const addBulkDownloadListeners = (startAppListening: AppStartListening) =
       });
     },
   });
-
-  startAppListening({
-    actionCreator: socketBulkDownloadStarted,
-    effect: async (action) => {
-      // This should always happen immediately after the bulk download request, so we don't need to show a toast here.
-      log.debug(action.payload.data, 'Bulk download preparation started');
-    },
-  });
-
-  startAppListening({
-    actionCreator: socketBulkDownloadComplete,
-    effect: async (action) => {
-      log.debug(action.payload.data, 'Bulk download preparation completed');
-
-      const { bulk_download_item_name } = action.payload.data;
-
-      // TODO(psyche): This URL may break in in some environments (e.g. Nvidia workbench) but we need to test it first
-      const url = `/api/v1/images/download/${bulk_download_item_name}`;
-
-      toast({
-        id: bulk_download_item_name,
-        title: t('gallery.bulkDownloadReady', 'Download ready'),
-        status: 'success',
-        description: (
-          <ExternalLink
-            label={t('gallery.clickToDownload', 'Click here to download')}
-            href={url}
-            download={bulk_download_item_name}
-          />
-        ),
-        duration: null,
-      });
-    },
-  });
-
-  startAppListening({
-    actionCreator: socketBulkDownloadError,
-    effect: async (action) => {
-      log.debug(action.payload.data, 'Bulk download preparation failed');
-
-      const { bulk_download_item_name } = action.payload.data;
-
-      toast({
-        id: bulk_download_item_name,
-        title: t('gallery.bulkDownloadFailed'),
-        status: 'error',
-        description: action.payload.data.error,
-        duration: null,
-      });
-    },
-  });
 };

invokeai/frontend/web/src/app/store/middleware/listenerMiddleware/listeners/cancellationsListeners.ts

@@ -0,0 +1,137 @@
+import type { AppStartListening } from 'app/store/middleware/listenerMiddleware';
+import { $lastCanvasProgressEvent } from 'features/controlLayers/store/canvasSlice';
+import { queueApi } from 'services/api/endpoints/queue';
+
+/**
+ * To prevent a race condition where a progress event arrives after a successful cancellation, we need to keep track of
+ * cancellations:
+ * - In the route handlers above, we track and update the cancellations object
+ * - When the user queues a, we should reset the cancellations, also handled int he route handlers above
+ * - When we get a progress event, we should check if the event is cancelled before setting the event
+ *
+ * We have a few ways that cancellations are effected, so we need to track them all:
+ * - by queue item id (in this case, we will compare the session_id and not the item_id)
+ * - by batch id
+ * - by destination
+ * - by clearing the queue
+ */
+type Cancellations = {
+  sessionIds: Set<string>;
+  batchIds: Set<string>;
+  destinations: Set<string>;
+  clearQueue: boolean;
+};
+
+const resetCancellations = (): void => {
+  cancellations.clearQueue = false;
+  cancellations.sessionIds.clear();
+  cancellations.batchIds.clear();
+  cancellations.destinations.clear();
+};
+
+const cancellations: Cancellations = {
+  sessionIds: new Set(),
+  batchIds: new Set(),
+  destinations: new Set(),
+  clearQueue: false,
+} as Readonly<Cancellations>;
+
+/**
+ * Checks if an item is cancelled, used to prevent race conditions with event handling.
+ *
+ * To use this, provide the session_id, batch_id and destination from the event payload.
+ */
+export const getIsCancelled = (item: {
+  session_id: string;
+  batch_id: string;
+  destination?: string | null;
+}): boolean => {
+  if (cancellations.clearQueue) {
+    return true;
+  }
+  if (cancellations.sessionIds.has(item.session_id)) {
+    return true;
+  }
+  if (cancellations.batchIds.has(item.batch_id)) {
+    return true;
+  }
+  if (item.destination && cancellations.destinations.has(item.destination)) {
+    return true;
+  }
+  return false;
+};
+
+export const addCancellationsListeners = (startAppListening: AppStartListening) => {
+  // When we get a cancellation, we may need to clear the last progress event - next few listeners handle those cases.
+  // Maybe we could use the `getIsCancelled` util here, but I think that could introduce _another_ race condition...
+  startAppListening({
+    matcher: queueApi.endpoints.enqueueBatch.matchFulfilled,
+    effect: () => {
+      resetCancellations();
+    },
+  });
+
+  startAppListening({
+    matcher: queueApi.endpoints.cancelByBatchDestination.matchFulfilled,
+    effect: (action) => {
+      cancellations.destinations.add(action.meta.arg.originalArgs.destination);
+
+      const event = $lastCanvasProgressEvent.get();
+      if (!event) {
+        return;
+      }
+      const { session_id, batch_id, destination } = event;
+      if (getIsCancelled({ session_id, batch_id, destination })) {
+        $lastCanvasProgressEvent.set(null);
+      }
+    },
+  });
+
+  startAppListening({
+    matcher: queueApi.endpoints.cancelQueueItem.matchFulfilled,
+    effect: (action) => {
+      cancellations.sessionIds.add(action.payload.session_id);
+
+      const event = $lastCanvasProgressEvent.get();
+      if (!event) {
+        return;
+      }
+      const { session_id, batch_id, destination } = event;
+      if (getIsCancelled({ session_id, batch_id, destination })) {
+        $lastCanvasProgressEvent.set(null);
+      }
+    },
+  });
+
+  startAppListening({
+    matcher: queueApi.endpoints.cancelByBatchIds.matchFulfilled,
+    effect: (action) => {
+      for (const batch_id of action.meta.arg.originalArgs.batch_ids) {
+        cancellations.batchIds.add(batch_id);
+      }
+      const event = $lastCanvasProgressEvent.get();
+      if (!event) {
+        return;
+      }
+      const { session_id, batch_id, destination } = event;
+      if (getIsCancelled({ session_id, batch_id, destination })) {
+        $lastCanvasProgressEvent.set(null);
+      }
+    },
+  });
+
+  startAppListening({
+    matcher: queueApi.endpoints.clearQueue.matchFulfilled,
+    effect: () => {
+      cancellations.clearQueue = true;
+      const event = $lastCanvasProgressEvent.get();
+      if (!event) {
+        return;
+      }
+      const { session_id, batch_id, destination } = event;
+      if (getIsCancelled({ session_id, batch_id, destination })) {
+        $lastCanvasProgressEvent.set(null);
+      }
+    },
+  });
+};

invokeai/frontend/web/src/app/store/middleware/listenerMiddleware/listeners/canvasCopiedToClipboard.ts

@@ -1,38 +0,0 @@
-import { $logger } from 'app/logging/logger';
-import type { AppStartListening } from 'app/store/middleware/listenerMiddleware';
-import { canvasCopiedToClipboard } from 'features/canvas/store/actions';
-import { getBaseLayerBlob } from 'features/canvas/util/getBaseLayerBlob';
-import { copyBlobToClipboard } from 'features/system/util/copyBlobToClipboard';
-import { toast } from 'features/toast/toast';
-import { t } from 'i18next';
-
-export const addCanvasCopiedToClipboardListener = (startAppListening: AppStartListening) => {
-  startAppListening({
-    actionCreator: canvasCopiedToClipboard,
-    effect: async (action, { getState }) => {
-      const moduleLog = $logger.get().child({ namespace: 'canvasCopiedToClipboardListener' });
-      const state = getState();
-
-      try {
-        const blob = getBaseLayerBlob(state);
-
-        copyBlobToClipboard(blob);
-      } catch (err) {
-        moduleLog.error(String(err));
-        toast({
-          id: 'CANVAS_COPY_FAILED',
-          title: t('toast.problemCopyingCanvas'),
-          description: t('toast.problemCopyingCanvasDesc'),
-          status: 'error',
-        });
-        return;
-      }
-
-      toast({
-        id: 'CANVAS_COPY_SUCCEEDED',
-        title: t('toast.canvasCopiedClipboard'),
-        status: 'success',
-      });
-    },
-  });
-};

invokeai/frontend/web/src/app/store/middleware/listenerMiddleware/listeners/canvasDownloadedAsImage.ts

@@ -1,34 +0,0 @@
-import { $logger } from 'app/logging/logger';
-import type { AppStartListening } from 'app/store/middleware/listenerMiddleware';
-import { canvasDownloadedAsImage } from 'features/canvas/store/actions';
-import { downloadBlob } from 'features/canvas/util/downloadBlob';
-import { getBaseLayerBlob } from 'features/canvas/util/getBaseLayerBlob';
-import { toast } from 'features/toast/toast';
-import { t } from 'i18next';
-
-export const addCanvasDownloadedAsImageListener = (startAppListening: AppStartListening) => {
-  startAppListening({
-    actionCreator: canvasDownloadedAsImage,
-    effect: async (action, { getState }) => {
-      const moduleLog = $logger.get().child({ namespace: 'canvasSavedToGalleryListener' });
-      const state = getState();
-
-      let blob;
-      try {
-        blob = await getBaseLayerBlob(state);
-      } catch (err) {
-        moduleLog.error(String(err));
-        toast({
-          id: 'CANVAS_DOWNLOAD_FAILED',
-          title: t('toast.problemDownloadingCanvas'),
-          description: t('toast.problemDownloadingCanvasDesc'),
-          status: 'error',
-        });
-        return;
-      }
-
-      downloadBlob(blob, 'canvas.png');
-      toast({ id: 'CANVAS_DOWNLOAD_SUCCEEDED', title: t('toast.canvasDownloaded'), status: 'success' });
-    },
-  });
-};

invokeai/frontend/web/src/app/store/middleware/listenerMiddleware/listeners/canvasImageToControlNet.ts

@@ -1,60 +0,0 @@
-import { logger } from 'app/logging/logger';
-import type { AppStartListening } from 'app/store/middleware/listenerMiddleware';
-import { canvasImageToControlAdapter } from 'features/canvas/store/actions';
-import { getBaseLayerBlob } from 'features/canvas/util/getBaseLayerBlob';
-import { controlAdapterImageChanged } from 'features/controlAdapters/store/controlAdaptersSlice';
-import { toast } from 'features/toast/toast';
-import { t } from 'i18next';
-import { imagesApi } from 'services/api/endpoints/images';
-
-export const addCanvasImageToControlNetListener = (startAppListening: AppStartListening) => {
-  startAppListening({
-    actionCreator: canvasImageToControlAdapter,
-    effect: async (action, { dispatch, getState }) => {
-      const log = logger('canvas');
-      const state = getState();
-      const { id } = action.payload;
-
-      let blob: Blob;
-      try {
-        blob = await getBaseLayerBlob(state, true);
-      } catch (err) {
-        log.error(String(err));
-        toast({
-          id: 'PROBLEM_SAVING_CANVAS',
-          title: t('toast.problemSavingCanvas'),
-          description: t('toast.problemSavingCanvasDesc'),
-          status: 'error',
-        });
-        return;
-      }
-
-      const { autoAddBoardId } = state.gallery;
-
-      const imageDTO = await dispatch(
-        imagesApi.endpoints.uploadImage.initiate({
-          file: new File([blob], 'savedCanvas.png', {
-            type: 'image/png',
-          }),
-          image_category: 'control',
-          is_intermediate: true,
-          board_id: autoAddBoardId === 'none' ? undefined : autoAddBoardId,
-          crop_visible: false,
-          postUploadAction: {
-            type: 'TOAST',
-            title: t('toast.canvasSentControlnetAssets'),
-          },
-        })
-      ).unwrap();
-
-      const { image_name } = imageDTO;
-
-      dispatch(
-        controlAdapterImageChanged({
-          id,
-          controlImage: image_name,
-        })
-      );
-    },
-  });
-};

invokeai/frontend/web/src/app/store/middleware/listenerMiddleware/listeners/canvasMaskSavedToGallery.ts

@@ -1,60 +0,0 @@
-import { logger } from 'app/logging/logger';
-import type { AppStartListening } from 'app/store/middleware/listenerMiddleware';
-import { canvasMaskSavedToGallery } from 'features/canvas/store/actions';
-import { getCanvasData } from 'features/canvas/util/getCanvasData';
-import { toast } from 'features/toast/toast';
-import { t } from 'i18next';
-import { imagesApi } from 'services/api/endpoints/images';
-
-export const addCanvasMaskSavedToGalleryListener = (startAppListening: AppStartListening) => {
-  startAppListening({
-    actionCreator: canvasMaskSavedToGallery,
-    effect: async (action, { dispatch, getState }) => {
-      const log = logger('canvas');
-      const state = getState();
-
-      const canvasBlobsAndImageData = await getCanvasData(
-        state.canvas.layerState,
-        state.canvas.boundingBoxCoordinates,
-        state.canvas.boundingBoxDimensions,
-        state.canvas.isMaskEnabled,
-        state.canvas.shouldPreserveMaskedArea
-      );
-
-      if (!canvasBlobsAndImageData) {
-        return;
-      }
-
-      const { maskBlob } = canvasBlobsAndImageData;
-
-      if (!maskBlob) {
-        log.error('Problem getting mask layer blob');
-        toast({
-          id: 'PROBLEM_SAVING_MASK',
-          title: t('toast.problemSavingMask'),
-          description: t('toast.problemSavingMaskDesc'),
-          status: 'error',
-        });
-        return;
-      }
-
-      const { autoAddBoardId } = state.gallery;
-
-      dispatch(
-        imagesApi.endpoints.uploadImage.initiate({
-          file: new File([maskBlob], 'canvasMaskImage.png', {
-            type: 'image/png',
-          }),
-          image_category: 'mask',
-          is_intermediate: false,
-          board_id: autoAddBoardId === 'none' ? undefined : autoAddBoardId,
-          crop_visible: true,
-          postUploadAction: {
-            type: 'TOAST',
-            title: t('toast.maskSavedAssets'),
-          },
-        })
-      );
-    },
-  });
-};

invokeai/frontend/web/src/app/store/middleware/listenerMiddleware/listeners/canvasMaskToControlNet.ts

@@ -1,70 +0,0 @@
-import { logger } from 'app/logging/logger';
-import type { AppStartListening } from 'app/store/middleware/listenerMiddleware';
-import { canvasMaskToControlAdapter } from 'features/canvas/store/actions';
-import { getCanvasData } from 'features/canvas/util/getCanvasData';
-import { controlAdapterImageChanged } from 'features/controlAdapters/store/controlAdaptersSlice';
-import { toast } from 'features/toast/toast';
-import { t } from 'i18next';
-import { imagesApi } from 'services/api/endpoints/images';
-
-export const addCanvasMaskToControlNetListener = (startAppListening: AppStartListening) => {
-  startAppListening({
-    actionCreator: canvasMaskToControlAdapter,
-    effect: async (action, { dispatch, getState }) => {
-      const log = logger('canvas');
-      const state = getState();
-      const { id } = action.payload;
-      const canvasBlobsAndImageData = await getCanvasData(
-        state.canvas.layerState,
-        state.canvas.boundingBoxCoordinates,
-        state.canvas.boundingBoxDimensions,
-        state.canvas.isMaskEnabled,
-        state.canvas.shouldPreserveMaskedArea
-      );
-
-      if (!canvasBlobsAndImageData) {
-        return;
-      }
-
-      const { maskBlob } = canvasBlobsAndImageData;
-
-      if (!maskBlob) {
-        log.error('Problem getting mask layer blob');
-        toast({
-          id: 'PROBLEM_IMPORTING_MASK',
-          title: t('toast.problemImportingMask'),
-          description: t('toast.problemImportingMaskDesc'),
-          status: 'error',
-        });
-        return;
-      }
-
-      const { autoAddBoardId } = state.gallery;
-
-      const imageDTO = await dispatch(
-        imagesApi.endpoints.uploadImage.initiate({
-          file: new File([maskBlob], 'canvasMaskImage.png', {
-            type: 'image/png',
-          }),
-          image_category: 'mask',
-          is_intermediate: true,
-          board_id: autoAddBoardId === 'none' ? undefined : autoAddBoardId,
-          crop_visible: false,
-          postUploadAction: {
-            type: 'TOAST',
-            title: t('toast.maskSentControlnetAssets'),
-          },
-        })
-      ).unwrap();
-
-      const { image_name } = imageDTO;
-
-      dispatch(
-        controlAdapterImageChanged({
-          id,
-          controlImage: image_name,
-        })
-      );
-    },
-  });
-};

invokeai/frontend/web/src/app/store/middleware/listenerMiddleware/listeners/canvasMerged.ts

@@ -1,73 +0,0 @@
-import { $logger } from 'app/logging/logger';
-import type { AppStartListening } from 'app/store/middleware/listenerMiddleware';
-import { canvasMerged } from 'features/canvas/store/actions';
-import { $canvasBaseLayer } from 'features/canvas/store/canvasNanostore';
-import { setMergedCanvas } from 'features/canvas/store/canvasSlice';
-import { getFullBaseLayerBlob } from 'features/canvas/util/getFullBaseLayerBlob';
-import { toast } from 'features/toast/toast';
-import { t } from 'i18next';
-import { imagesApi } from 'services/api/endpoints/images';
-
-export const addCanvasMergedListener = (startAppListening: AppStartListening) => {
-  startAppListening({
-    actionCreator: canvasMerged,
-    effect: async (action, { dispatch }) => {
-      const moduleLog = $logger.get().child({ namespace: 'canvasCopiedToClipboardListener' });
-      const blob = await getFullBaseLayerBlob();
-
-      if (!blob) {
-        moduleLog.error('Problem getting base layer blob');
-        toast({
-          id: 'PROBLEM_MERGING_CANVAS',
-          title: t('toast.problemMergingCanvas'),
-          description: t('toast.problemMergingCanvasDesc'),
-          status: 'error',
-        });
-        return;
-      }
-
-      const canvasBaseLayer = $canvasBaseLayer.get();
-
-      if (!canvasBaseLayer) {
-        moduleLog.error('Problem getting canvas base layer');
-        toast({
-          id: 'PROBLEM_MERGING_CANVAS',
-          title: t('toast.problemMergingCanvas'),
-          description: t('toast.problemMergingCanvasDesc'),
-          status: 'error',
-        });
-        return;
-      }
-
-      const baseLayerRect = canvasBaseLayer.getClientRect({
-        relativeTo: canvasBaseLayer.getParent() ?? undefined,
-      });
-
-      const imageDTO = await dispatch(
-        imagesApi.endpoints.uploadImage.initiate({
-          file: new File([blob], 'mergedCanvas.png', {
-            type: 'image/png',
-          }),
-          image_category: 'general',
-          is_intermediate: true,
-          postUploadAction: {
-            type: 'TOAST',
-            title: t('toast.canvasMerged'),
-          },
-        })
-      ).unwrap();
-
-      // TODO: I can't figure out how to do the type narrowing in the `take()` so just brute forcing it here
-      const { image_name } = imageDTO;
-
-      dispatch(
-        setMergedCanvas({
-          kind: 'image',
-          layer: 'base',
-          imageName: image_name,
-          ...baseLayerRect,
-        })
-      );
-    },
-  });
-};

invokeai/frontend/web/src/app/store/middleware/listenerMiddleware/listeners/canvasSavedToGallery.ts

@@ -1,53 +0,0 @@
-import { logger } from 'app/logging/logger';
-import type { AppStartListening } from 'app/store/middleware/listenerMiddleware';
-import { parseify } from 'common/util/serialize';
-import { canvasSavedToGallery } from 'features/canvas/store/actions';
-import { getBaseLayerBlob } from 'features/canvas/util/getBaseLayerBlob';
-import { toast } from 'features/toast/toast';
-import { t } from 'i18next';
-import { imagesApi } from 'services/api/endpoints/images';
-
-export const addCanvasSavedToGalleryListener = (startAppListening: AppStartListening) => {
-  startAppListening({
-    actionCreator: canvasSavedToGallery,
-    effect: async (action, { dispatch, getState }) => {
-      const log = logger('canvas');
-      const state = getState();
-
-      let blob;
-      try {
-        blob = await getBaseLayerBlob(state);
-      } catch (err) {
-        log.error(String(err));
-        toast({
-          id: 'CANVAS_SAVE_FAILED',
-          title: t('toast.problemSavingCanvas'),
-          description: t('toast.problemSavingCanvasDesc'),
-          status: 'error',
-        });
-        return;
-      }
-
-      const { autoAddBoardId } = state.gallery;
-
-      dispatch(
-        imagesApi.endpoints.uploadImage.initiate({
-          file: new File([blob], 'savedCanvas.png', {
-            type: 'image/png',
-          }),
-          image_category: 'general',
-          is_intermediate: false,
-          board_id: autoAddBoardId === 'none' ? undefined : autoAddBoardId,
-          crop_visible: true,
-          postUploadAction: {
-            type: 'TOAST',
-            title: t('toast.canvasSavedGallery'),
-          },
-          metadata: {
-            _canvas_objects: parseify(state.canvas.layerState.objects),
-          },
-        })
-      );
-    },
-  });
-};

invokeai/frontend/web/src/app/store/middleware/listenerMiddleware/listeners/controlAdapterPreprocessor.ts

@@ -1,194 +0,0 @@
-import { isAnyOf } from '@reduxjs/toolkit';
-import { logger } from 'app/logging/logger';
-import type { AppStartListening } from 'app/store/middleware/listenerMiddleware';
-import type { AppDispatch } from 'app/store/store';
-import { parseify } from 'common/util/serialize';
-import {
-  caLayerImageChanged,
-  caLayerModelChanged,
-  caLayerProcessedImageChanged,
-  caLayerProcessorConfigChanged,
-  caLayerProcessorPendingBatchIdChanged,
-  caLayerRecalled,
-  isControlAdapterLayer,
-} from 'features/controlLayers/store/controlLayersSlice';
-import { CA_PROCESSOR_DATA } from 'features/controlLayers/util/controlAdapters';
-import { toast } from 'features/toast/toast';
-import { t } from 'i18next';
-import { isEqual } from 'lodash-es';
-import { getImageDTO } from 'services/api/endpoints/images';
-import { queueApi } from 'services/api/endpoints/queue';
-import type { BatchConfig } from 'services/api/types';
-import { socketInvocationComplete } from 'services/events/actions';
-import { assert } from 'tsafe';
-
-const matcher = isAnyOf(
-  caLayerImageChanged,
-  caLayerProcessedImageChanged,
-  caLayerProcessorConfigChanged,
-  caLayerModelChanged,
-  caLayerRecalled
-);
-
-const DEBOUNCE_MS = 300;
-const log = logger('session');
-
-/**
- * Simple helper to cancel a batch and reset the pending batch ID
- */
-const cancelProcessorBatch = async (dispatch: AppDispatch, layerId: string, batchId: string) => {
-  const req = dispatch(queueApi.endpoints.cancelByBatchIds.initiate({ batch_ids: [batchId] }));
-  log.trace({ batchId }, 'Cancelling existing preprocessor batch');
-  try {
-    await req.unwrap();
-  } catch {
-    // no-op
-  } finally {
-    req.reset();
-    // Always reset the pending batch ID - the cancel req could fail if the batch doesn't exist
-    dispatch(caLayerProcessorPendingBatchIdChanged({ layerId, batchId: null }));
-  }
-};
-
-export const addControlAdapterPreprocessor = (startAppListening: AppStartListening) => {
-  startAppListening({
-    matcher,
-    effect: async (action, { dispatch, getState, getOriginalState, cancelActiveListeners, delay, take, signal }) => {
-      const layerId = caLayerRecalled.match(action) ? action.payload.id : action.payload.layerId;
-      const state = getState();
-      const originalState = getOriginalState();
-
-      // Cancel any in-progress instances of this listener
-      cancelActiveListeners();
-      log.trace('Control Layer CA auto-process triggered');
-
-      // Delay before starting actual work
-      await delay(DEBOUNCE_MS);
-
-      const layer = state.controlLayers.present.layers.filter(isControlAdapterLayer).find((l) => l.id === layerId);
-
-      if (!layer) {
-        return;
-      }
-
-      // We should only process if the processor settings or image have changed
-      const originalLayer = originalState.controlLayers.present.layers
-        .filter(isControlAdapterLayer)
-        .find((l) => l.id === layerId);
-      const originalImage = originalLayer?.controlAdapter.image;
-      const originalConfig = originalLayer?.controlAdapter.processorConfig;
-
-      const image = layer.controlAdapter.image;
-      const processedImage = layer.controlAdapter.processedImage;
-      const config = layer.controlAdapter.processorConfig;
-
-      if (isEqual(config, originalConfig) && isEqual(image, originalImage) && processedImage) {
-        // Neither config nor image have changed, we can bail
-        return;
-      }
-
-      if (!image || !config) {
-        // - If we have no image, we have nothing to process
-        // - If we have no processor config, we have nothing to process
-        // Clear the processed image and bail
-        dispatch(caLayerProcessedImageChanged({ layerId, imageDTO: null }));
-        return;
-      }
-
-      // At this point, the user has stopped fiddling with the processor settings and there is a processor selected.
-
-      // If there is a pending processor batch, cancel it.
-      if (layer.controlAdapter.processorPendingBatchId) {
-        cancelProcessorBatch(dispatch, layerId, layer.controlAdapter.processorPendingBatchId);
-      }
-
-      // TODO(psyche): I can't get TS to be happy, it thinkgs `config` is `never` but it should be inferred from the generic... I'll just cast it for now
-      const processorNode = CA_PROCESSOR_DATA[config.type].buildNode(image, config as never);
-      const enqueueBatchArg: BatchConfig = {
-        prepend: true,
-        batch: {
-          graph: {
-            nodes: {
-              [processorNode.id]: {
-                ...processorNode,
-                // Control images are always intermediate - do not save to gallery
-                is_intermediate: true,
-              },
-            },
-            edges: [],
-          },
-          runs: 1,
-        },
-      };
-
-      // Kick off the processor batch
-      const req = dispatch(
-        queueApi.endpoints.enqueueBatch.initiate(enqueueBatchArg, {
-          fixedCacheKey: 'enqueueBatch',
-        })
-      );
-
-      try {
-        const enqueueResult = await req.unwrap();
-        // TODO(psyche): Update the pydantic models, pretty sure we will _always_ have a batch_id here, but the model says it's optional
-        assert(enqueueResult.batch.batch_id, 'Batch ID not returned from queue');
-        dispatch(caLayerProcessorPendingBatchIdChanged({ layerId, batchId: enqueueResult.batch.batch_id }));
-        log.debug({ enqueueResult: parseify(enqueueResult) }, t('queue.graphQueued'));
-
-        // Wait for the processor node to complete
-        const [invocationCompleteAction] = await take(
-          (action): action is ReturnType<typeof socketInvocationComplete> =>
-            socketInvocationComplete.match(action) &&
-            action.payload.data.batch_id === enqueueResult.batch.batch_id &&
-            action.payload.data.invocation_source_id === processorNode.id
-        );
-
-        // We still have to check the output type
-        assert(
-          invocationCompleteAction.payload.data.result.type === 'image_output',
-          `Processor did not return an image output, got: ${invocationCompleteAction.payload.data.result}`
-        );
-        const { image_name } = invocationCompleteAction.payload.data.result.image;
-
-        const imageDTO = await getImageDTO(image_name);
-        assert(imageDTO, "Failed to fetch processor output's image DTO");
-
-        // Whew! We made it. Update the layer with the processed image
-        log.debug({ layerId, imageDTO }, 'ControlNet image processed');
-        dispatch(caLayerProcessedImageChanged({ layerId, imageDTO }));
-        dispatch(caLayerProcessorPendingBatchIdChanged({ layerId, batchId: null }));
-      } catch (error) {
-        if (signal.aborted) {
-          // The listener was canceled - we need to cancel the pending processor batch, if there is one (could have changed by now).
-          const pendingBatchId = getState()
-            .controlLayers.present.layers.filter(isControlAdapterLayer)
-            .find((l) => l.id === layerId)?.controlAdapter.processorPendingBatchId;
-          if (pendingBatchId) {
-            cancelProcessorBatch(dispatch, layerId, pendingBatchId);
-          }
-          log.trace('Control Adapter preprocessor cancelled');
-        } else {
-          // Some other error condition...
-          log.error({ enqueueBatchArg: parseify(enqueueBatchArg) }, t('queue.graphFailedToQueue'));
-
-          if (error instanceof Object) {
-            if ('data' in error && 'status' in error) {
-              if (error.status === 403) {
-                dispatch(caLayerImageChanged({ layerId, imageDTO: null }));
-                return;
-              }
-            }
-          }
-
-          toast({
-            id: 'GRAPH_QUEUE_FAILED',
-            title: t('queue.graphFailedToQueue'),
-            status: 'error',
-          });
-        }
-      } finally {
-        req.reset();
-      }
-    },
-  });
-};

invokeai/frontend/web/src/app/store/middleware/listenerMiddleware/listeners/controlNetAutoProcess.ts

@@ -1,85 +0,0 @@
-import type { AnyListenerPredicate } from '@reduxjs/toolkit';
-import { logger } from 'app/logging/logger';
-import type { AppStartListening } from 'app/store/middleware/listenerMiddleware';
-import type { RootState } from 'app/store/store';
-import { controlAdapterImageProcessed } from 'features/controlAdapters/store/actions';
-import {
-  controlAdapterAutoConfigToggled,
-  controlAdapterImageChanged,
-  controlAdapterModelChanged,
-  controlAdapterProcessorParamsChanged,
-  controlAdapterProcessortTypeChanged,
-  selectControlAdapterById,
-} from 'features/controlAdapters/store/controlAdaptersSlice';
-import { isControlNetOrT2IAdapter } from 'features/controlAdapters/store/types';
-
-type AnyControlAdapterParamChangeAction =
-  | ReturnType<typeof controlAdapterProcessorParamsChanged>
-  | ReturnType<typeof controlAdapterModelChanged>
-  | ReturnType<typeof controlAdapterImageChanged>
-  | ReturnType<typeof controlAdapterProcessortTypeChanged>
-  | ReturnType<typeof controlAdapterAutoConfigToggled>;
-
-const predicate: AnyListenerPredicate<RootState> = (action, state, prevState) => {
-  const isActionMatched =
-    controlAdapterProcessorParamsChanged.match(action) ||
-    controlAdapterModelChanged.match(action) ||
-    controlAdapterImageChanged.match(action) ||
-    controlAdapterProcessortTypeChanged.match(action) ||
-    controlAdapterAutoConfigToggled.match(action);
-
-  if (!isActionMatched) {
-    return false;
-  }
-
-  const { id } = action.payload;
-  const prevCA = selectControlAdapterById(prevState.controlAdapters, id);
-  const ca = selectControlAdapterById(state.controlAdapters, id);
-  if (!prevCA || !isControlNetOrT2IAdapter(prevCA) || !ca || !isControlNetOrT2IAdapter(ca)) {
-    return false;
-  }
-
-  if (controlAdapterAutoConfigToggled.match(action)) {
-    // do not process if the user just disabled auto-config
-    if (prevCA.shouldAutoConfig === true) {
-      return false;
-    }
-  }
-
-  const { controlImage, processorType, shouldAutoConfig } = ca;
-  if (controlAdapterModelChanged.match(action) && !shouldAutoConfig) {
-    // do not process if the action is a model change but the processor settings are dirty
-    return false;
-  }
-
-  const isProcessorSelected = processorType !== 'none';
-
-  const hasControlImage = Boolean(controlImage);
-
-  return isProcessorSelected && hasControlImage;
-};
-
-const DEBOUNCE_MS = 300;
-
-/**
- * Listener that automatically processes a ControlNet image when its processor parameters are changed.
- *
- * The network request is debounced.
- */
-export const addControlNetAutoProcessListener = (startAppListening: AppStartListening) => {
-  startAppListening({
-    predicate,
-    effect: async (action, { dispatch, cancelActiveListeners, delay }) => {
-      const log = logger('session');
-      const { id } = (action as AnyControlAdapterParamChangeAction).payload;
-
-      // Cancel any in-progress instances of this listener
-      cancelActiveListeners();
-      log.trace('ControlNet auto-process triggered');
-      // Delay before starting actual work
-      await delay(DEBOUNCE_MS);
-
-      dispatch(controlAdapterImageProcessed({ id }));
-    },
-  });
-};

invokeai/frontend/web/src/app/store/middleware/listenerMiddleware/listeners/controlNetImageProcessed.ts

@@ -1,118 +0,0 @@
-import { logger } from 'app/logging/logger';
-import type { AppStartListening } from 'app/store/middleware/listenerMiddleware';
-import { parseify } from 'common/util/serialize';
-import { controlAdapterImageProcessed } from 'features/controlAdapters/store/actions';
-import {
-  controlAdapterImageChanged,
-  controlAdapterProcessedImageChanged,
-  pendingControlImagesCleared,
-  selectControlAdapterById,
-} from 'features/controlAdapters/store/controlAdaptersSlice';
-import { isControlNetOrT2IAdapter } from 'features/controlAdapters/store/types';
-import { toast } from 'features/toast/toast';
-import { t } from 'i18next';
-import { imagesApi } from 'services/api/endpoints/images';
-import { queueApi } from 'services/api/endpoints/queue';
-import type { BatchConfig, ImageDTO } from 'services/api/types';
-import { socketInvocationComplete } from 'services/events/actions';
-
-export const addControlNetImageProcessedListener = (startAppListening: AppStartListening) => {
-  startAppListening({
-    actionCreator: controlAdapterImageProcessed,
-    effect: async (action, { dispatch, getState, take }) => {
-      const log = logger('session');
-      const { id } = action.payload;
-      const ca = selectControlAdapterById(getState().controlAdapters, id);
-
-      if (!ca?.controlImage || !isControlNetOrT2IAdapter(ca)) {
-        log.error('Unable to process ControlNet image');
-        return;
-      }
-
-      if (ca.processorType === 'none' || ca.processorNode.type === 'none') {
-        return;
-      }
-
-      // ControlNet one-off procressing graph is just the processor node, no edges.
-      // Also we need to grab the image.
-
-      const nodeId = ca.processorNode.id;
-      const enqueueBatchArg: BatchConfig = {
-        prepend: true,
-        batch: {
-          graph: {
-            nodes: {
-              [ca.processorNode.id]: {
-                ...ca.processorNode,
-                is_intermediate: true,
-                use_cache: false,
-                image: { image_name: ca.controlImage },
-              },
-            },
-            edges: [],
-          },
-          runs: 1,
-        },
-      };
-
-      try {
-        const req = dispatch(
-          queueApi.endpoints.enqueueBatch.initiate(enqueueBatchArg, {
-            fixedCacheKey: 'enqueueBatch',
-          })
-        );
-        const enqueueResult = await req.unwrap();
-        req.reset();
-        log.debug({ enqueueResult: parseify(enqueueResult) }, t('queue.graphQueued'));
-
-        const [invocationCompleteAction] = await take(
-          (action): action is ReturnType<typeof socketInvocationComplete> =>
-            socketInvocationComplete.match(action) &&
-            action.payload.data.batch_id === enqueueResult.batch.batch_id &&
-            action.payload.data.invocation_source_id === nodeId
-        );
-
-        // We still have to check the output type
-        if (invocationCompleteAction.payload.data.result.type === 'image_output') {
-          const { image_name } = invocationCompleteAction.payload.data.result.image;
-
-          // Wait for the ImageDTO to be received
-          const [{ payload }] = await take(
-            (action) =>
-              imagesApi.endpoints.getImageDTO.matchFulfilled(action) && action.payload.image_name === image_name
-          );
-
-          const processedControlImage = payload as ImageDTO;
-
-          log.debug({ controlNetId: action.payload, processedControlImage }, 'ControlNet image processed');
-
-          // Update the processed image in the store
-          dispatch(
-            controlAdapterProcessedImageChanged({
-              id,
-              processedControlImage: processedControlImage.image_name,
-            })
-          );
-        }
-      } catch (error) {
-        log.error({ enqueueBatchArg: parseify(enqueueBatchArg) }, t('queue.graphFailedToQueue'));
-
-        if (error instanceof Object) {
-          if ('data' in error && 'status' in error) {
-            if (error.status === 403) {
-              dispatch(pendingControlImagesCleared());
-              dispatch(controlAdapterImageChanged({ id, controlImage: null }));
-              return;
-            }
-          }
-        }
-
-        toast({
-          id: 'GRAPH_QUEUE_FAILED',
-          title: t('queue.graphFailedToQueue'),
-          status: 'error',
-        });
-      }
-    },
-  });
-};

invokeai/frontend/web/src/app/store/middleware/listenerMiddleware/listeners/enqueueRequestedCanvas.ts

@@ -1,144 +0,0 @@
-import { logger } from 'app/logging/logger';
-import { enqueueRequested } from 'app/store/actions';
-import type { AppStartListening } from 'app/store/middleware/listenerMiddleware';
-import openBase64ImageInTab from 'common/util/openBase64ImageInTab';
-import { parseify } from 'common/util/serialize';
-import { canvasBatchIdAdded, stagingAreaInitialized } from 'features/canvas/store/canvasSlice';
-import { blobToDataURL } from 'features/canvas/util/blobToDataURL';
-import { getCanvasData } from 'features/canvas/util/getCanvasData';
-import { getCanvasGenerationMode } from 'features/canvas/util/getCanvasGenerationMode';
-import { canvasGraphBuilt } from 'features/nodes/store/actions';
-import { prepareLinearUIBatch } from 'features/nodes/util/graph/buildLinearBatchConfig';
-import { buildCanvasGraph } from 'features/nodes/util/graph/canvas/buildCanvasGraph';
-import { imagesApi } from 'services/api/endpoints/images';
-import { queueApi } from 'services/api/endpoints/queue';
-import type { ImageDTO } from 'services/api/types';
-
-/**
- * This listener is responsible invoking the canvas. This involves a number of steps:
- *
- * 1. Generate image blobs from the canvas layers
- * 2. Determine the generation mode from the layers (txt2img, img2img, inpaint)
- * 3. Build the canvas graph
- * 4. Create the session with the graph
- * 5. Upload the init image if necessary
- * 6. Upload the mask image if necessary
- * 7. Update the init and mask images with the session ID
- * 8. Initialize the staging area if not yet initialized
- * 9. Dispatch the sessionReadyToInvoke action to invoke the session
- */
-export const addEnqueueRequestedCanvasListener = (startAppListening: AppStartListening) => {
-  startAppListening({
-    predicate: (action): action is ReturnType<typeof enqueueRequested> =>
-      enqueueRequested.match(action) && action.payload.tabName === 'canvas',
-    effect: async (action, { getState, dispatch }) => {
-      const log = logger('queue');
-      const { prepend } = action.payload;
-      const state = getState();
-
-      const { layerState, boundingBoxCoordinates, boundingBoxDimensions, isMaskEnabled, shouldPreserveMaskedArea } =
-        state.canvas;
-
-      // Build canvas blobs
-      const canvasBlobsAndImageData = await getCanvasData(
-        layerState,
-        boundingBoxCoordinates,
-        boundingBoxDimensions,
-        isMaskEnabled,
-        shouldPreserveMaskedArea
-      );
-
-      if (!canvasBlobsAndImageData) {
-        log.error('Unable to create canvas data');
-        return;
-      }
-
-      const { baseBlob, baseImageData, maskBlob, maskImageData } = canvasBlobsAndImageData;
-
-      // Determine the generation mode
-      const generationMode = getCanvasGenerationMode(baseImageData, maskImageData);
-
-      if (state.system.enableImageDebugging) {
-        const baseDataURL = await blobToDataURL(baseBlob);
-        const maskDataURL = await blobToDataURL(maskBlob);
-        openBase64ImageInTab([
-          { base64: maskDataURL, caption: 'mask b64' },
-          { base64: baseDataURL, caption: 'image b64' },
-        ]);
-      }
-
-      log.debug(`Generation mode: ${generationMode}`);
-
-      // Temp placeholders for the init and mask images
-      let canvasInitImage: ImageDTO | undefined;
-      let canvasMaskImage: ImageDTO | undefined;
-
-      // For img2img and inpaint/outpaint, we need to upload the init images
-      if (['img2img', 'inpaint', 'outpaint'].includes(generationMode)) {
-        // upload the image, saving the request id
-        canvasInitImage = await dispatch(
-          imagesApi.endpoints.uploadImage.initiate({
-            file: new File([baseBlob], 'canvasInitImage.png', {
-              type: 'image/png',
-            }),
-            image_category: 'general',
-            is_intermediate: true,
-          })
-        ).unwrap();
-      }
-
-      // For inpaint/outpaint, we also need to upload the mask layer
-      if (['inpaint', 'outpaint'].includes(generationMode)) {
-        // upload the image, saving the request id
-        canvasMaskImage = await dispatch(
-          imagesApi.endpoints.uploadImage.initiate({
-            file: new File([maskBlob], 'canvasMaskImage.png', {
-              type: 'image/png',
-            }),
-            image_category: 'mask',
-            is_intermediate: true,
-          })
-        ).unwrap();
-      }
-
-      const graph = await buildCanvasGraph(state, generationMode, canvasInitImage, canvasMaskImage);
-
-      log.debug({ graph: parseify(graph) }, `Canvas graph built`);
-
-      // currently this action is just listened to for logging
-      dispatch(canvasGraphBuilt(graph));
-
-      const batchConfig = prepareLinearUIBatch(state, graph, prepend);
-
-      try {
-        const req = dispatch(
-          queueApi.endpoints.enqueueBatch.initiate(batchConfig, {
-            fixedCacheKey: 'enqueueBatch',
-          })
-        );
-
-        const enqueueResult = await req.unwrap();
-        req.reset();
-
-        const batchId = enqueueResult.batch.batch_id as string; // we know the is a string, backend provides it
-
-        // Prep the canvas staging area if it is not yet initialized
-        if (!state.canvas.layerState.stagingArea.boundingBox) {
-          dispatch(
-            stagingAreaInitialized({
-              boundingBox: {
-                ...state.canvas.boundingBoxCoordinates,
-                ...state.canvas.boundingBoxDimensions,
-              },
-            })
-          );
-        }
-
-        // Associate the session with the canvas session ID
-        dispatch(canvasBatchIdAdded(batchId));
-      } catch {
-        // no-op
-      }
-    },
-  });
-};

invokeai/frontend/web/src/app/store/middleware/listenerMiddleware/listeners/enqueueRequestedLinear.ts

@@ -1,10 +1,25 @@
+import { logger } from 'app/logging/logger';
 import { enqueueRequested } from 'app/store/actions';
 import type { AppStartListening } from 'app/store/middleware/listenerMiddleware';
-import { isImageViewerOpenChanged } from 'features/gallery/store/gallerySlice';
+import type { SerializableObject } from 'common/types';
+import type { Result } from 'common/util/result';
+import { isErr, withResult, withResultAsync } from 'common/util/result';
+import { $canvasManager } from 'features/controlLayers/store/canvasSlice';
+import {
+  selectIsStaging,
+  stagingAreaReset,
+  stagingAreaStartedStaging,
+} from 'features/controlLayers/store/canvasStagingAreaSlice';
 import { prepareLinearUIBatch } from 'features/nodes/util/graph/buildLinearBatchConfig';
-import { buildGenerationTabGraph } from 'features/nodes/util/graph/generation/buildGenerationTabGraph';
-import { buildGenerationTabSDXLGraph } from 'features/nodes/util/graph/generation/buildGenerationTabSDXLGraph';
+import { buildSD1Graph } from 'features/nodes/util/graph/generation/buildSD1Graph';
+import { buildSDXLGraph } from 'features/nodes/util/graph/generation/buildSDXLGraph';
+import type { Graph } from 'features/nodes/util/graph/generation/Graph';
+import { serializeError } from 'serialize-error';
 import { queueApi } from 'services/api/endpoints/queue';
+import type { Invocation } from 'services/api/types';
+import { assert } from 'tsafe';
+
+const log = logger('generation');
 
 export const addEnqueueRequestedLinear = (startAppListening: AppStartListening) => {
   startAppListening({
@@ -12,33 +27,81 @@ export const addEnqueueRequestedLinear = (startAppListening: AppStartListening)
       enqueueRequested.match(action) && action.payload.tabName === 'generation',
     effect: async (action, { getState, dispatch }) => {
       const state = getState();
-      const { shouldShowProgressInViewer } = state.ui;
-      const model = state.generation.model;
+      const model = state.params.model;
       const { prepend } = action.payload;
 
-      let graph;
+      const manager = $canvasManager.get();
+      assert(manager, 'No model found in state');
 
-      if (model?.base === 'sdxl') {
-        graph = await buildGenerationTabSDXLGraph(state);
-      } else {
-        graph = await buildGenerationTabGraph(state);
+      let didStartStaging = false;
+
+      if (!selectIsStaging(state) && state.canvasSettings.sendToCanvas) {
+        dispatch(stagingAreaStartedStaging());
+        didStartStaging = true;
       }
 
-      const batchConfig = prepareLinearUIBatch(state, graph, prepend);
+      const abortStaging = () => {
+        if (didStartStaging && selectIsStaging(getState())) {
+          dispatch(stagingAreaReset());
+        }
+      };
+
+      let buildGraphResult: Result<
+        { g: Graph; noise: Invocation<'noise'>; posCond: Invocation<'compel' | 'sdxl_compel_prompt'> },
+        Error
+      >;
+
+      assert(model, 'No model found in state');
+      const base = model.base;
+
+      switch (base) {
+        case 'sdxl':
+          buildGraphResult = await withResultAsync(() => buildSDXLGraph(state, manager));
+          break;
+        case 'sd-1':
+        case `sd-2`:
+          buildGraphResult = await withResultAsync(() => buildSD1Graph(state, manager));
+          break;
+        default:
+          assert(false, `No graph builders for base ${base}`);
+      }
+
+      if (isErr(buildGraphResult)) {
+        log.error({ error: serializeError(buildGraphResult.error) }, 'Failed to build graph');
+        abortStaging();
+        return;
+      }
+
+      const { g, noise, posCond } = buildGraphResult.value;
+
+      const destination = state.canvasSettings.sendToCanvas ? 'canvas' : 'gallery';
+
+      const prepareBatchResult = withResult(() =>
+        prepareLinearUIBatch(state, g, prepend, noise, posCond, 'generation', destination)
+      );
+
+      if (isErr(prepareBatchResult)) {
+        log.error({ error: serializeError(prepareBatchResult.error) }, 'Failed to prepare batch');
+        abortStaging();
+        return;
+      }
 
       const req = dispatch(
-        queueApi.endpoints.enqueueBatch.initiate(batchConfig, {
+        queueApi.endpoints.enqueueBatch.initiate(prepareBatchResult.value, {
           fixedCacheKey: 'enqueueBatch',
         })
       );
-      try {
-        await req.unwrap();
-        if (shouldShowProgressInViewer) {
-          dispatch(isImageViewerOpenChanged(true));
-        }
-      } finally {
-        req.reset();
+      req.reset();
+
+      const enqueueResult = await withResultAsync(() => req.unwrap());
+
+      if (isErr(enqueueResult)) {
+        log.error({ error: serializeError(enqueueResult.error) }, 'Failed to enqueue batch');
+        abortStaging();
+        return;
       }
+
+      log.debug({ batchConfig: prepareBatchResult.value } as SerializableObject, 'Enqueued batch');
     },
   });
 };

invokeai/frontend/web/src/app/store/middleware/listenerMiddleware/listeners/enqueueRequestedNodes.ts

@@ -1,5 +1,6 @@
 import { enqueueRequested } from 'app/store/actions';
 import type { AppStartListening } from 'app/store/middleware/listenerMiddleware';
+import { selectNodesSlice } from 'features/nodes/store/selectors';
 import { buildNodesGraph } from 'features/nodes/util/graph/buildNodesGraph';
 import { buildWorkflowWithValidation } from 'features/nodes/util/workflow/buildWorkflow';
 import { queueApi } from 'services/api/endpoints/queue';
@@ -11,12 +12,12 @@ export const addEnqueueRequestedNodes = (startAppListening: AppStartListening) =
       enqueueRequested.match(action) && action.payload.tabName === 'workflows',
     effect: async (action, { getState, dispatch }) => {
       const state = getState();
-      const { nodes, edges } = state.nodes.present;
+      const nodes = selectNodesSlice(state);
       const workflow = state.workflow;
-      const graph = buildNodesGraph(state.nodes.present);
+      const graph = buildNodesGraph(nodes);
       const builtWorkflow = buildWorkflowWithValidation({
-        nodes,
-        edges,
+        nodes: nodes.nodes,
+        edges: nodes.edges,
         workflow,
       });
 
@@ -29,7 +30,9 @@ export const addEnqueueRequestedNodes = (startAppListening: AppStartListening) =
         batch: {
           graph,
           workflow: builtWorkflow,
-          runs: state.generation.iterations,
+          runs: state.params.iterations,
+          origin: 'workflows',
+          destination: 'gallery',
         },
         prepend: action.payload.prepend,
       };

invokeai/frontend/web/src/app/store/middleware/listenerMiddleware/listeners/enqueueRequestedUpscale.ts

@@ -1,6 +1,5 @@
 import { enqueueRequested } from 'app/store/actions';
 import type { AppStartListening } from 'app/store/middleware/listenerMiddleware';
-import { isImageViewerOpenChanged } from 'features/gallery/store/gallerySlice';
 import { prepareLinearUIBatch } from 'features/nodes/util/graph/buildLinearBatchConfig';
 import { buildMultidiffusionUpscaleGraph } from 'features/nodes/util/graph/buildMultidiffusionUpscaleGraph';
 import { queueApi } from 'services/api/endpoints/queue';
@@ -11,12 +10,11 @@ export const addEnqueueRequestedUpscale = (startAppListening: AppStartListening)
       enqueueRequested.match(action) && action.payload.tabName === 'upscaling',
     effect: async (action, { getState, dispatch }) => {
       const state = getState();
-      const { shouldShowProgressInViewer } = state.ui;
       const { prepend } = action.payload;
 
-      const graph = await buildMultidiffusionUpscaleGraph(state);
+      const { g, noise, posCond } = await buildMultidiffusionUpscaleGraph(state);
 
-      const batchConfig = prepareLinearUIBatch(state, graph, prepend);
+      const batchConfig = prepareLinearUIBatch(state, g, prepend, noise, posCond, 'upscaling', 'gallery');
 
       const req = dispatch(
         queueApi.endpoints.enqueueBatch.initiate(batchConfig, {
@@ -25,9 +23,6 @@ export const addEnqueueRequestedUpscale = (startAppListening: AppStartListening)
       );
       try {
         await req.unwrap();
-        if (shouldShowProgressInViewer) {
-          dispatch(isImageViewerOpenChanged(true));
-        }
       } finally {
         req.reset();
       }

invokeai/frontend/web/src/app/store/middleware/listenerMiddleware/listeners/galleryImageClicked.ts

@@ -27,7 +27,7 @@ export const galleryImageClicked = createAction<{
 export const addGalleryImageClickedListener = (startAppListening: AppStartListening) => {
   startAppListening({
     actionCreator: galleryImageClicked,
-    effect: async (action, { dispatch, getState }) => {
+    effect: (action, { dispatch, getState }) => {
       const { imageDTO, shiftKey, ctrlKey, metaKey, altKey } = action.payload;
       const state = getState();
       const queryArgs = selectListImagesQueryArgs(state);

invokeai/frontend/web/src/app/store/middleware/listenerMiddleware/listeners/getOpenAPISchema.ts

@@ -1,24 +1,27 @@
 import { logger } from 'app/logging/logger';
 import type { AppStartListening } from 'app/store/middleware/listenerMiddleware';
+import type { SerializableObject } from 'common/types';
 import { parseify } from 'common/util/serialize';
 import { $templates } from 'features/nodes/store/nodesSlice';
 import { parseSchema } from 'features/nodes/util/schema/parseSchema';
 import { size } from 'lodash-es';
+import { serializeError } from 'serialize-error';
 import { appInfoApi } from 'services/api/endpoints/appInfo';
 
+const log = logger('system');
+
 export const addGetOpenAPISchemaListener = (startAppListening: AppStartListening) => {
   startAppListening({
     matcher: appInfoApi.endpoints.getOpenAPISchema.matchFulfilled,
     effect: (action, { getState }) => {
-      const log = logger('system');
       const schemaJSON = action.payload;
 
-      log.debug({ schemaJSON: parseify(schemaJSON) }, 'Received OpenAPI schema');
+      log.debug({ schemaJSON: parseify(schemaJSON) } as SerializableObject, 'Received OpenAPI schema');
       const { nodesAllowlist, nodesDenylist } = getState().config;
 
       const nodeTemplates = parseSchema(schemaJSON, nodesAllowlist, nodesDenylist);
 
-      log.debug({ nodeTemplates: parseify(nodeTemplates) }, `Built ${size(nodeTemplates)} node templates`);
+      log.debug({ nodeTemplates } as SerializableObject, `Built ${size(nodeTemplates)} node templates`);
 
       $templates.set(nodeTemplates);
     },
@@ -30,8 +33,7 @@ export const addGetOpenAPISchemaListener = (startAppListening: AppStartListening
       // If action.meta.condition === true, the request was canceled/skipped because another request was in flight or
       // the value was already in the cache. We don't want to log these errors.
       if (!action.meta.condition) {
-        const log = logger('system');
-        log.error({ error: parseify(action.error) }, 'Problem retrieving OpenAPI Schema');
+        log.error({ error: serializeError(action.error) }, 'Problem retrieving OpenAPI Schema');
       }
     },
   });

invokeai/frontend/web/src/app/store/middleware/listenerMiddleware/listeners/imageAddedToBoard.ts

@@ -2,15 +2,13 @@ import { logger } from 'app/logging/logger';
 import type { AppStartListening } from 'app/store/middleware/listenerMiddleware';
 import { imagesApi } from 'services/api/endpoints/images';
 
+const log = logger('gallery');
+
 export const addImageAddedToBoardFulfilledListener = (startAppListening: AppStartListening) => {
   startAppListening({
     matcher: imagesApi.endpoints.addImageToBoard.matchFulfilled,
     effect: (action) => {
-      const log = logger('images');
       const { board_id, imageDTO } = action.meta.arg.originalArgs;
-
-      // TODO: update listImages cache for this board
-
       log.debug({ board_id, imageDTO }, 'Image added to board');
     },
   });
@@ -18,9 +16,7 @@ export const addImageAddedToBoardFulfilledListener = (startAppListening: AppStar
   startAppListening({
     matcher: imagesApi.endpoints.addImageToBoard.matchRejected,
     effect: (action) => {
-      const log = logger('images');
       const { board_id, imageDTO } = action.meta.arg.originalArgs;
-
       log.debug({ board_id, imageDTO }, 'Problem adding image to board');
     },
   });

invokeai/frontend/web/src/app/store/middleware/listenerMiddleware/listeners/imageDeletionListeners.ts

@@ -1,20 +1,9 @@
 import { logger } from 'app/logging/logger';
 import type { AppStartListening } from 'app/store/middleware/listenerMiddleware';
 import type { AppDispatch, RootState } from 'app/store/store';
-import { resetCanvas } from 'features/canvas/store/canvasSlice';
-import {
-  controlAdapterImageChanged,
-  controlAdapterProcessedImageChanged,
-  selectControlAdapterAll,
-} from 'features/controlAdapters/store/controlAdaptersSlice';
-import { isControlNetOrT2IAdapter } from 'features/controlAdapters/store/types';
-import {
-  isControlAdapterLayer,
-  isInitialImageLayer,
-  isIPAdapterLayer,
-  isRegionalGuidanceLayer,
-  layerDeleted,
-} from 'features/controlLayers/store/controlLayersSlice';
+import { entityDeleted, ipaImageChanged } from 'features/controlLayers/store/canvasSlice';
+import { selectCanvasSlice } from 'features/controlLayers/store/selectors';
+import { getEntityIdentifier } from 'features/controlLayers/store/types';
 import { imageDeletionConfirmed } from 'features/deleteImageModal/store/actions';
 import { isModalOpenChanged } from 'features/deleteImageModal/store/slice';
 import { selectListImagesQueryArgs } from 'features/gallery/store/gallerySelectors';
@@ -26,6 +15,10 @@ import { forEach, intersectionBy } from 'lodash-es';
 import { imagesApi } from 'services/api/endpoints/images';
 import type { ImageDTO } from 'services/api/types';
 
+const log = logger('gallery');
+
+//TODO(psyche): handle image deletion (canvas staging area?)
+
 // Some utils to delete images from different parts of the app
 const deleteNodesImages = (state: RootState, dispatch: AppDispatch, imageDTO: ImageDTO) => {
   state.nodes.present.nodes.forEach((node) => {
@@ -47,52 +40,37 @@ const deleteNodesImages = (state: RootState, dispatch: AppDispatch, imageDTO: Im
   });
 };
 
-const deleteControlAdapterImages = (state: RootState, dispatch: AppDispatch, imageDTO: ImageDTO) => {
-  forEach(selectControlAdapterAll(state.controlAdapters), (ca) => {
-    if (
-      ca.controlImage === imageDTO.image_name ||
-      (isControlNetOrT2IAdapter(ca) && ca.processedControlImage === imageDTO.image_name)
-    ) {
-      dispatch(
-        controlAdapterImageChanged({
-          id: ca.id,
-          controlImage: null,
-        })
-      );
-      dispatch(
-        controlAdapterProcessedImageChanged({
-          id: ca.id,
-          processedControlImage: null,
-        })
-      );
+// const deleteControlAdapterImages = (state: RootState, dispatch: AppDispatch, imageDTO: ImageDTO) => {
+//   state.canvas.present.controlAdapters.entities.forEach(({ id, imageObject, processedImageObject }) => {
+//     if (
+//       imageObject?.image.image_name === imageDTO.image_name ||
+//       processedImageObject?.image.image_name === imageDTO.image_name
+//     ) {
+//       dispatch(caImageChanged({ id, imageDTO: null }));
+//       dispatch(caProcessedImageChanged({ id, imageDTO: null }));
+//     }
+//   });
+// };
+
+const deleteIPAdapterImages = (state: RootState, dispatch: AppDispatch, imageDTO: ImageDTO) => {
+  selectCanvasSlice(state).ipAdapters.entities.forEach((entity) => {
+    if (entity.ipAdapter.image?.image_name === imageDTO.image_name) {
+      dispatch(ipaImageChanged({ entityIdentifier: getEntityIdentifier(entity), imageDTO: null }));
     }
   });
 };
 
-const deleteControlLayerImages = (state: RootState, dispatch: AppDispatch, imageDTO: ImageDTO) => {
-  state.controlLayers.present.layers.forEach((l) => {
-    if (isRegionalGuidanceLayer(l)) {
-      if (l.ipAdapters.some((ipa) => ipa.image?.name === imageDTO.image_name)) {
-        dispatch(layerDeleted(l.id));
+const deleteLayerImages = (state: RootState, dispatch: AppDispatch, imageDTO: ImageDTO) => {
+  selectCanvasSlice(state).rasterLayers.entities.forEach(({ id, objects }) => {
+    let shouldDelete = false;
+    for (const obj of objects) {
+      if (obj.type === 'image' && obj.image.image_name === imageDTO.image_name) {
+        shouldDelete = true;
+        break;
       }
     }
-    if (isControlAdapterLayer(l)) {
-      if (
-        l.controlAdapter.image?.name === imageDTO.image_name ||
-        l.controlAdapter.processedImage?.name === imageDTO.image_name
-      ) {
-        dispatch(layerDeleted(l.id));
-      }
-    }
-    if (isIPAdapterLayer(l)) {
-      if (l.ipAdapter.image?.name === imageDTO.image_name) {
-        dispatch(layerDeleted(l.id));
-      }
-    }
-    if (isInitialImageLayer(l)) {
-      if (l.image?.name === imageDTO.image_name) {
-        dispatch(layerDeleted(l.id));
-      }
+    if (shouldDelete) {
+      dispatch(entityDeleted({ entityIdentifier: { id, type: 'raster_layer' } }));
     }
   });
 };
@@ -145,14 +123,10 @@ export const addImageDeletionListeners = (startAppListening: AppStartListening)
           }
         }
 
-        // We need to reset the features where the image is in use - none of these work if their image(s) don't exist
-        if (imageUsage.isCanvasImage) {
-          dispatch(resetCanvas());
-        }
-
-        deleteControlAdapterImages(state, dispatch, imageDTO);
         deleteNodesImages(state, dispatch, imageDTO);
-        deleteControlLayerImages(state, dispatch, imageDTO);
+        // deleteControlAdapterImages(state, dispatch, imageDTO);
+        deleteIPAdapterImages(state, dispatch, imageDTO);
+        deleteLayerImages(state, dispatch, imageDTO);
       } catch {
         // no-op
       } finally {
@@ -189,14 +163,11 @@ export const addImageDeletionListeners = (startAppListening: AppStartListening)
 
         // We need to reset the features where the image is in use - none of these work if their image(s) don't exist
 
-        if (imagesUsage.some((i) => i.isCanvasImage)) {
-          dispatch(resetCanvas());
-        }
-
         imageDTOs.forEach((imageDTO) => {
-          deleteControlAdapterImages(state, dispatch, imageDTO);
           deleteNodesImages(state, dispatch, imageDTO);
-          deleteControlLayerImages(state, dispatch, imageDTO);
+          // deleteControlAdapterImages(state, dispatch, imageDTO);
+          deleteIPAdapterImages(state, dispatch, imageDTO);
+          deleteLayerImages(state, dispatch, imageDTO);
         });
       } catch {
         // no-op
@@ -220,7 +191,6 @@ export const addImageDeletionListeners = (startAppListening: AppStartListening)
   startAppListening({
     matcher: imagesApi.endpoints.deleteImage.matchFulfilled,
     effect: (action) => {
-      const log = logger('images');
       log.debug({ imageDTO: action.meta.arg.originalArgs }, 'Image deleted');
     },
   });
@@ -228,7 +198,6 @@ export const addImageDeletionListeners = (startAppListening: AppStartListening)
   startAppListening({
     matcher: imagesApi.endpoints.deleteImage.matchRejected,
     effect: (action) => {
-      const log = logger('images');
       log.debug({ imageDTO: action.meta.arg.originalArgs }, 'Unable to delete image');
     },
   });

invokeai/frontend/web/src/app/store/middleware/listenerMiddleware/listeners/imageDropped.ts

@@ -1,28 +1,20 @@
 import { createAction } from '@reduxjs/toolkit';
 import { logger } from 'app/logging/logger';
 import type { AppStartListening } from 'app/store/middleware/listenerMiddleware';
-import { parseify } from 'common/util/serialize';
-import { setInitialCanvasImage } from 'features/canvas/store/canvasSlice';
+import { selectDefaultControlAdapter } from 'features/controlLayers/hooks/addLayerHooks';
 import {
-  controlAdapterImageChanged,
-  controlAdapterIsEnabledChanged,
-} from 'features/controlAdapters/store/controlAdaptersSlice';
-import {
-  caLayerImageChanged,
-  iiLayerImageChanged,
-  ipaLayerImageChanged,
-  rgLayerIPAdapterImageChanged,
-} from 'features/controlLayers/store/controlLayersSlice';
+  controlLayerAdded,
+  ipaImageChanged,
+  rasterLayerAdded,
+  rgIPAdapterImageChanged,
+} from 'features/controlLayers/store/canvasSlice';
+import { selectCanvasSlice } from 'features/controlLayers/store/selectors';
+import type { CanvasControlLayerState, CanvasRasterLayerState } from 'features/controlLayers/store/types';
+import { imageDTOToImageObject } from 'features/controlLayers/store/types';
 import type { TypesafeDraggableData, TypesafeDroppableData } from 'features/dnd/types';
 import { isValidDrop } from 'features/dnd/util/isValidDrop';
-import {
-  imageSelected,
-  imageToCompareChanged,
-  isImageViewerOpenChanged,
-  selectionChanged,
-} from 'features/gallery/store/gallerySlice';
+import { imageToCompareChanged, selectionChanged } from 'features/gallery/store/gallerySlice';
 import { fieldImageValueChanged } from 'features/nodes/store/nodesSlice';
-import { selectOptimalDimension } from 'features/parameters/store/generationSlice';
 import { upscaleInitialImageChanged } from 'features/parameters/store/upscaleSlice';
 import { imagesApi } from 'services/api/endpoints/images';
 
@@ -31,11 +23,12 @@ export const dndDropped = createAction<{
   activeData: TypesafeDraggableData;
 }>('dnd/dndDropped');
 
+const log = logger('system');
+
 export const addImageDroppedListener = (startAppListening: AppStartListening) => {
   startAppListening({
     actionCreator: dndDropped,
-    effect: async (action, { dispatch, getState }) => {
-      const log = logger('dnd');
+    effect: (action, { dispatch, getState }) => {
       const { activeData, overData } = action.payload;
       if (!isValidDrop(overData, activeData)) {
         return;
@@ -46,80 +39,22 @@ export const addImageDroppedListener = (startAppListening: AppStartListening) =>
       } else if (activeData.payloadType === 'GALLERY_SELECTION') {
         log.debug({ activeData, overData }, `Images (${getState().gallery.selection.length}) dropped`);
       } else if (activeData.payloadType === 'NODE_FIELD') {
-        log.debug({ activeData: parseify(activeData), overData: parseify(overData) }, 'Node field dropped');
+        log.debug({ activeData, overData }, 'Node field dropped');
       } else {
         log.debug({ activeData, overData }, `Unknown payload dropped`);
       }
 
-      /**
-       * Image dropped on current image
-       */
-      if (
-        overData.actionType === 'SET_CURRENT_IMAGE' &&
-        activeData.payloadType === 'IMAGE_DTO' &&
-        activeData.payload.imageDTO
-      ) {
-        dispatch(imageSelected(activeData.payload.imageDTO));
-        dispatch(isImageViewerOpenChanged(true));
-        return;
-      }
-
-      /**
-       * Image dropped on ControlNet
-       */
-      if (
-        overData.actionType === 'SET_CONTROL_ADAPTER_IMAGE' &&
-        activeData.payloadType === 'IMAGE_DTO' &&
-        activeData.payload.imageDTO
-      ) {
-        const { id } = overData.context;
-        dispatch(
-          controlAdapterImageChanged({
-            id,
-            controlImage: activeData.payload.imageDTO.image_name,
-          })
-        );
-        dispatch(
-          controlAdapterIsEnabledChanged({
-            id,
-            isEnabled: true,
-          })
-        );
-        return;
-      }
-
-      /**
-       * Image dropped on Control Adapter Layer
-       */
-      if (
-        overData.actionType === 'SET_CA_LAYER_IMAGE' &&
-        activeData.payloadType === 'IMAGE_DTO' &&
-        activeData.payload.imageDTO
-      ) {
-        const { layerId } = overData.context;
-        dispatch(
-          caLayerImageChanged({
-            layerId,
-            imageDTO: activeData.payload.imageDTO,
-          })
-        );
-        return;
-      }
-
       /**
        * Image dropped on IP Adapter Layer
        */
       if (
-        overData.actionType === 'SET_IPA_LAYER_IMAGE' &&
+        overData.actionType === 'SET_IPA_IMAGE' &&
         activeData.payloadType === 'IMAGE_DTO' &&
         activeData.payload.imageDTO
       ) {
-        const { layerId } = overData.context;
+        const { id } = overData.context;
         dispatch(
-          ipaLayerImageChanged({
-            layerId,
-            imageDTO: activeData.payload.imageDTO,
-          })
+          ipaImageChanged({ entityIdentifier: { id, type: 'ip_adapter' }, imageDTO: activeData.payload.imageDTO })
         );
         return;
       }
@@ -128,14 +63,14 @@ export const addImageDroppedListener = (startAppListening: AppStartListening) =>
        * Image dropped on RG Layer IP Adapter
        */
       if (
-        overData.actionType === 'SET_RG_LAYER_IP_ADAPTER_IMAGE' &&
+        overData.actionType === 'SET_RG_IP_ADAPTER_IMAGE' &&
         activeData.payloadType === 'IMAGE_DTO' &&
         activeData.payload.imageDTO
       ) {
-        const { layerId, ipAdapterId } = overData.context;
+        const { id, ipAdapterId } = overData.context;
         dispatch(
-          rgLayerIPAdapterImageChanged({
-            layerId,
+          rgIPAdapterImageChanged({
+            entityIdentifier: { id, type: 'regional_guidance' },
             ipAdapterId,
             imageDTO: activeData.payload.imageDTO,
           })
@@ -144,32 +79,41 @@ export const addImageDroppedListener = (startAppListening: AppStartListening) =>
       }
 
       /**
-       * Image dropped on II Layer Image
+       * Image dropped on Raster layer
        */
       if (
-        overData.actionType === 'SET_II_LAYER_IMAGE' &&
+        overData.actionType === 'ADD_RASTER_LAYER_FROM_IMAGE' &&
         activeData.payloadType === 'IMAGE_DTO' &&
         activeData.payload.imageDTO
       ) {
-        const { layerId } = overData.context;
-        dispatch(
-          iiLayerImageChanged({
-            layerId,
-            imageDTO: activeData.payload.imageDTO,
-          })
-        );
+        const imageObject = imageDTOToImageObject(activeData.payload.imageDTO);
+        const { x, y } = selectCanvasSlice(getState()).bbox.rect;
+        const overrides: Partial<CanvasRasterLayerState> = {
+          objects: [imageObject],
+          position: { x, y },
+        };
+        dispatch(rasterLayerAdded({ overrides, isSelected: true }));
         return;
       }
 
       /**
-       * Image dropped on Canvas
+       * Image dropped on Raster layer
        */
       if (
-        overData.actionType === 'SET_CANVAS_INITIAL_IMAGE' &&
+        overData.actionType === 'ADD_CONTROL_LAYER_FROM_IMAGE' &&
         activeData.payloadType === 'IMAGE_DTO' &&
         activeData.payload.imageDTO
       ) {
-        dispatch(setInitialCanvasImage(activeData.payload.imageDTO, selectOptimalDimension(getState())));
+        const state = getState();
+        const imageObject = imageDTOToImageObject(activeData.payload.imageDTO);
+        const { x, y } = selectCanvasSlice(state).bbox.rect;
+        const defaultControlAdapter = selectDefaultControlAdapter(state);
+        const overrides: Partial<CanvasControlLayerState> = {
+          objects: [imageObject],
+          position: { x, y },
+          controlAdapter: defaultControlAdapter,
+        };
+        dispatch(controlLayerAdded({ overrides, isSelected: true }));
         return;
       }
 
@@ -202,7 +146,6 @@ export const addImageDroppedListener = (startAppListening: AppStartListening) =>
       ) {
         const { imageDTO } = activeData.payload;
         dispatch(imageToCompareChanged(imageDTO));
-        dispatch(isImageViewerOpenChanged(true));
         return;
       }
 

invokeai/frontend/web/src/app/store/middleware/listenerMiddleware/listeners/imageRemovedFromBoard.ts

@@ -2,13 +2,13 @@ import { logger } from 'app/logging/logger';
 import type { AppStartListening } from 'app/store/middleware/listenerMiddleware';
 import { imagesApi } from 'services/api/endpoints/images';
 
+const log = logger('gallery');
+
 export const addImageRemovedFromBoardFulfilledListener = (startAppListening: AppStartListening) => {
   startAppListening({
     matcher: imagesApi.endpoints.removeImageFromBoard.matchFulfilled,
     effect: (action) => {
-      const log = logger('images');
       const imageDTO = action.meta.arg.originalArgs;
-
       log.debug({ imageDTO }, 'Image removed from board');
     },
   });
@@ -16,9 +16,7 @@ export const addImageRemovedFromBoardFulfilledListener = (startAppListening: App
   startAppListening({
     matcher: imagesApi.endpoints.removeImageFromBoard.matchRejected,
     effect: (action) => {
-      const log = logger('images');
       const imageDTO = action.meta.arg.originalArgs;
-
       log.debug({ imageDTO }, 'Problem removing image from board');
     },
   });

invokeai/frontend/web/src/app/store/middleware/listenerMiddleware/listeners/imageToDeleteSelected.ts

@@ -6,16 +6,17 @@ import { imagesToDeleteSelected, isModalOpenChanged } from 'features/deleteImage
 export const addImageToDeleteSelectedListener = (startAppListening: AppStartListening) => {
   startAppListening({
     actionCreator: imagesToDeleteSelected,
-    effect: async (action, { dispatch, getState }) => {
+    effect: (action, { dispatch, getState }) => {
       const imageDTOs = action.payload;
       const state = getState();
       const { shouldConfirmOnDelete } = state.system;
       const imagesUsage = selectImageUsage(getState());
 
       const isImageInUse =
-        imagesUsage.some((i) => i.isCanvasImage) ||
-        imagesUsage.some((i) => i.isControlImage) ||
-        imagesUsage.some((i) => i.isNodesImage);
+        imagesUsage.some((i) => i.isLayerImage) ||
+        imagesUsage.some((i) => i.isControlAdapterImage) ||
+        imagesUsage.some((i) => i.isIPAdapterImage) ||
+        imagesUsage.some((i) => i.isLayerImage);
 
       if (shouldConfirmOnDelete || isImageInUse) {
         dispatch(isModalOpenChanged(true));

invokeai/frontend/web/src/app/store/middleware/listenerMiddleware/listeners/imageUploaded.ts

@@ -1,19 +1,9 @@
 import { logger } from 'app/logging/logger';
 import type { AppStartListening } from 'app/store/middleware/listenerMiddleware';
-import { setInitialCanvasImage } from 'features/canvas/store/canvasSlice';
-import {
-  controlAdapterImageChanged,
-  controlAdapterIsEnabledChanged,
-} from 'features/controlAdapters/store/controlAdaptersSlice';
-import {
-  caLayerImageChanged,
-  iiLayerImageChanged,
-  ipaLayerImageChanged,
-  rgLayerIPAdapterImageChanged,
-} from 'features/controlLayers/store/controlLayersSlice';
+import { ipaImageChanged, rgIPAdapterImageChanged } from 'features/controlLayers/store/canvasSlice';
 import { selectListBoardsQueryArgs } from 'features/gallery/store/gallerySelectors';
+import { boardIdSelected, galleryViewChanged } from 'features/gallery/store/gallerySlice';
 import { fieldImageValueChanged } from 'features/nodes/store/nodesSlice';
-import { selectOptimalDimension } from 'features/parameters/store/generationSlice';
 import { upscaleInitialImageChanged } from 'features/parameters/store/upscaleSlice';
 import { toast } from 'features/toast/toast';
 import { t } from 'i18next';
@@ -21,11 +11,12 @@ import { omit } from 'lodash-es';
 import { boardsApi } from 'services/api/endpoints/boards';
 import { imagesApi } from 'services/api/endpoints/images';
 
+const log = logger('gallery');
+
 export const addImageUploadedFulfilledListener = (startAppListening: AppStartListening) => {
   startAppListening({
     matcher: imagesApi.endpoints.uploadImage.matchFulfilled,
     effect: (action, { dispatch, getState }) => {
-      const log = logger('images');
       const imageDTO = action.payload;
       const state = getState();
       const { autoAddBoardId } = state.gallery;
@@ -54,6 +45,8 @@ export const addImageUploadedFulfilledListener = (startAppListening: AppStartLis
         if (!autoAddBoardId || autoAddBoardId === 'none') {
           const title = postUploadAction.title || DEFAULT_UPLOADED_TOAST.title;
           toast({ ...DEFAULT_UPLOADED_TOAST, title });
+          dispatch(boardIdSelected({ boardId: 'none' }));
+          dispatch(galleryViewChanged('assets'));
         } else {
           // Add this image to the board
           dispatch(
@@ -77,19 +70,12 @@ export const addImageUploadedFulfilledListener = (startAppListening: AppStartLis
             ...DEFAULT_UPLOADED_TOAST,
             description,
           });
+          dispatch(boardIdSelected({ boardId: autoAddBoardId }));
+          dispatch(galleryViewChanged('assets'));
         }
         return;
       }
 
-      if (postUploadAction?.type === 'SET_CANVAS_INITIAL_IMAGE') {
-        dispatch(setInitialCanvasImage(imageDTO, selectOptimalDimension(state)));
-        toast({
-          ...DEFAULT_UPLOADED_TOAST,
-          description: t('toast.setAsCanvasInitialImage'),
-        });
-        return;
-      }
-
       if (postUploadAction?.type === 'SET_UPSCALE_INITIAL_IMAGE') {
         dispatch(upscaleInitialImageChanged(imageDTO));
         toast({
@@ -99,70 +85,33 @@ export const addImageUploadedFulfilledListener = (startAppListening: AppStartLis
         return;
       }
 
-      if (postUploadAction?.type === 'SET_CONTROL_ADAPTER_IMAGE') {
+      // if (postUploadAction?.type === 'SET_CA_IMAGE') {
+      //   const { id } = postUploadAction;
+      //   dispatch(caImageChanged({ id, imageDTO }));
+      //   toast({ ...DEFAULT_UPLOADED_TOAST, description: t('toast.setControlImage') });
+      //   return;
+      // }
+
+      if (postUploadAction?.type === 'SET_IPA_IMAGE') {
         const { id } = postUploadAction;
-        dispatch(
-          controlAdapterIsEnabledChanged({
-            id,
-            isEnabled: true,
-          })
-        );
-        dispatch(
-          controlAdapterImageChanged({
-            id,
-            controlImage: imageDTO.image_name,
-          })
-        );
-        toast({
-          ...DEFAULT_UPLOADED_TOAST,
-          description: t('toast.setControlImage'),
-        });
+        dispatch(ipaImageChanged({ entityIdentifier: { id, type: 'ip_adapter' }, imageDTO }));
+        toast({ ...DEFAULT_UPLOADED_TOAST, description: t('toast.setControlImage') });
         return;
       }
 
-      if (postUploadAction?.type === 'SET_CA_LAYER_IMAGE') {
-        const { layerId } = postUploadAction;
-        dispatch(caLayerImageChanged({ layerId, imageDTO }));
-        toast({
-          ...DEFAULT_UPLOADED_TOAST,
-          description: t('toast.setControlImage'),
-        });
-      }
-
-      if (postUploadAction?.type === 'SET_IPA_LAYER_IMAGE') {
-        const { layerId } = postUploadAction;
-        dispatch(ipaLayerImageChanged({ layerId, imageDTO }));
-        toast({
-          ...DEFAULT_UPLOADED_TOAST,
-          description: t('toast.setControlImage'),
-        });
-      }
-
-      if (postUploadAction?.type === 'SET_RG_LAYER_IP_ADAPTER_IMAGE') {
-        const { layerId, ipAdapterId } = postUploadAction;
-        dispatch(rgLayerIPAdapterImageChanged({ layerId, ipAdapterId, imageDTO }));
-        toast({
-          ...DEFAULT_UPLOADED_TOAST,
-          description: t('toast.setControlImage'),
-        });
-      }
-
-      if (postUploadAction?.type === 'SET_II_LAYER_IMAGE') {
-        const { layerId } = postUploadAction;
-        dispatch(iiLayerImageChanged({ layerId, imageDTO }));
-        toast({
-          ...DEFAULT_UPLOADED_TOAST,
-          description: t('toast.setControlImage'),
-        });
+      if (postUploadAction?.type === 'SET_RG_IP_ADAPTER_IMAGE') {
+        const { id, ipAdapterId } = postUploadAction;
+        dispatch(
+          rgIPAdapterImageChanged({ entityIdentifier: { id, type: 'regional_guidance' }, ipAdapterId, imageDTO })
+        );
+        toast({ ...DEFAULT_UPLOADED_TOAST, description: t('toast.setControlImage') });
+        return;
       }
 
       if (postUploadAction?.type === 'SET_NODES_IMAGE') {
         const { nodeId, fieldName } = postUploadAction;
         dispatch(fieldImageValueChanged({ nodeId, fieldName, value: imageDTO }));
-        toast({
-          ...DEFAULT_UPLOADED_TOAST,
-          description: `${t('toast.setNodeField')} ${fieldName}`,
-        });
+        toast({ ...DEFAULT_UPLOADED_TOAST, description: `${t('toast.setNodeField')} ${fieldName}` });
         return;
       }
     },
@@ -171,7 +120,6 @@ export const addImageUploadedFulfilledListener = (startAppListening: AppStartLis
   startAppListening({
     matcher: imagesApi.endpoints.uploadImage.matchRejected,
     effect: (action) => {
-      const log = logger('images');
       const sanitizedData = {
         arg: {
           ...omit(action.meta.arg.originalArgs, ['file', 'postUploadAction']),

invokeai/frontend/web/src/app/store/middleware/listenerMiddleware/listeners/imagesStarred.ts

@@ -6,7 +6,7 @@ import type { ImageDTO } from 'services/api/types';
 export const addImagesStarredListener = (startAppListening: AppStartListening) => {
   startAppListening({
     matcher: imagesApi.endpoints.starImages.matchFulfilled,
-    effect: async (action, { dispatch, getState }) => {
+    effect: (action, { dispatch, getState }) => {
       const { updated_image_names: starredImages } = action.payload;
 
       const state = getState();

invokeai/frontend/web/src/app/store/middleware/listenerMiddleware/listeners/imagesUnstarred.ts

@@ -6,7 +6,7 @@ import type { ImageDTO } from 'services/api/types';
 export const addImagesUnstarredListener = (startAppListening: AppStartListening) => {
   startAppListening({
     matcher: imagesApi.endpoints.unstarImages.matchFulfilled,
-    effect: async (action, { dispatch, getState }) => {
+    effect: (action, { dispatch, getState }) => {
       const { updated_image_names: unstarredImages } = action.payload;
 
       const state = getState();

invokeai/frontend/web/src/app/store/middleware/listenerMiddleware/listeners/modelSelected.ts

@@ -1,23 +1,18 @@
 import { logger } from 'app/logging/logger';
 import type { AppStartListening } from 'app/store/middleware/listenerMiddleware';
-import {
-  controlAdapterIsEnabledChanged,
-  selectControlAdapterAll,
-} from 'features/controlAdapters/store/controlAdaptersSlice';
-import { loraRemoved } from 'features/lora/store/loraSlice';
+import { loraDeleted } from 'features/controlLayers/store/lorasSlice';
+import { modelChanged, vaeSelected } from 'features/controlLayers/store/paramsSlice';
 import { modelSelected } from 'features/parameters/store/actions';
-import { modelChanged, vaeSelected } from 'features/parameters/store/generationSlice';
 import { zParameterModel } from 'features/parameters/types/parameterSchemas';
 import { toast } from 'features/toast/toast';
 import { t } from 'i18next';
-import { forEach } from 'lodash-es';
+
+const log = logger('models');
 
 export const addModelSelectedListener = (startAppListening: AppStartListening) => {
   startAppListening({
     actionCreator: modelSelected,
     effect: (action, { getState, dispatch }) => {
-      const log = logger('models');
-
       const state = getState();
       const result = zParameterModel.safeParse(action.payload);
 
@@ -29,34 +24,36 @@ export const addModelSelectedListener = (startAppListening: AppStartListening) =
       const newModel = result.data;
 
       const newBaseModel = newModel.base;
-      const didBaseModelChange = state.generation.model?.base !== newBaseModel;
+      const didBaseModelChange = state.params.model?.base !== newBaseModel;
 
       if (didBaseModelChange) {
         // we may need to reset some incompatible submodels
         let modelsCleared = 0;
 
         // handle incompatible loras
-        forEach(state.lora.loras, (lora, id) => {
+        state.loras.loras.forEach((lora) => {
           if (lora.model.base !== newBaseModel) {
-            dispatch(loraRemoved(id));
+            dispatch(loraDeleted({ id: lora.id }));
             modelsCleared += 1;
           }
         });
 
         // handle incompatible vae
-        const { vae } = state.generation;
+        const { vae } = state.params;
         if (vae && vae.base !== newBaseModel) {
           dispatch(vaeSelected(null));
           modelsCleared += 1;
         }
 
         // handle incompatible controlnets
-        selectControlAdapterAll(state.controlAdapters).forEach((ca) => {
-          if (ca.model?.base !== newBaseModel) {
-            dispatch(controlAdapterIsEnabledChanged({ id: ca.id, isEnabled: false }));
-            modelsCleared += 1;
-          }
-        });
+        // state.canvas.present.controlAdapters.entities.forEach((ca) => {
+        //   if (ca.model?.base !== newBaseModel) {
+        //     modelsCleared += 1;
+        //     if (ca.isEnabled) {
+        //       dispatch(entityIsEnabledToggled({ entityIdentifier: { id: ca.id, type: 'control_adapter' } }));
+        //     }
+        //   }
+        // });
 
         if (modelsCleared > 0) {
           toast({
@@ -70,7 +67,7 @@ export const addModelSelectedListener = (startAppListening: AppStartListening) =
         }
       }
 
-      dispatch(modelChanged(newModel, state.generation.model));
+      dispatch(modelChanged({ model: newModel, previousModel: state.params.model }));
     },
   });
 };

invokeai/frontend/web/src/app/store/middleware/listenerMiddleware/listeners/modelsLoaded.ts

@@ -1,36 +1,42 @@
 import { logger } from 'app/logging/logger';
 import type { AppStartListening } from 'app/store/middleware/listenerMiddleware';
 import type { AppDispatch, RootState } from 'app/store/store';
-import type { JSONObject } from 'common/types';
+import type { SerializableObject } from 'common/types';
 import {
-  controlAdapterModelCleared,
-  selectControlAdapterAll,
-} from 'features/controlAdapters/store/controlAdaptersSlice';
-import { heightChanged, widthChanged } from 'features/controlLayers/store/controlLayersSlice';
-import { loraRemoved } from 'features/lora/store/loraSlice';
-import { calculateNewSize } from 'features/parameters/components/ImageSize/calculateNewSize';
-import { modelChanged, vaeSelected } from 'features/parameters/store/generationSlice';
+  bboxHeightChanged,
+  bboxWidthChanged,
+  controlLayerModelChanged,
+  ipaModelChanged,
+  rgIPAdapterModelChanged,
+} from 'features/controlLayers/store/canvasSlice';
+import { loraDeleted } from 'features/controlLayers/store/lorasSlice';
+import { modelChanged, refinerModelChanged, vaeSelected } from 'features/controlLayers/store/paramsSlice';
+import { selectCanvasSlice } from 'features/controlLayers/store/selectors';
+import { getEntityIdentifier } from 'features/controlLayers/store/types';
+import { calculateNewSize } from 'features/parameters/components/Bbox/calculateNewSize';
 import { postProcessingModelChanged, upscaleModelChanged } from 'features/parameters/store/upscaleSlice';
 import { zParameterModel, zParameterVAEModel } from 'features/parameters/types/parameterSchemas';
 import { getIsSizeOptimal, getOptimalDimension } from 'features/parameters/util/optimalDimension';
-import { refinerModelChanged } from 'features/sdxl/store/sdxlSlice';
-import { forEach } from 'lodash-es';
 import type { Logger } from 'roarr';
 import { modelConfigsAdapterSelectors, modelsApi } from 'services/api/endpoints/models';
 import type { AnyModelConfig } from 'services/api/types';
 import {
+  isControlNetOrT2IAdapterModelConfig,
+  isIPAdapterModelConfig,
+  isLoRAModelConfig,
   isNonRefinerMainModelConfig,
   isRefinerMainModelModelConfig,
   isSpandrelImageToImageModelConfig,
   isVAEModelConfig,
 } from 'services/api/types';
 
+const log = logger('models');
+
 export const addModelsLoadedListener = (startAppListening: AppStartListening) => {
   startAppListening({
     predicate: modelsApi.endpoints.getModelConfigs.matchFulfilled,
-    effect: async (action, { getState, dispatch }) => {
+    effect: (action, { getState, dispatch }) => {
       // models loaded, we need to ensure the selected model is available and if not, select the first one
-      const log = logger('models');
       log.info({ models: action.payload.entities }, `Models loaded (${action.payload.ids.length})`);
 
       const state = getState();
@@ -43,6 +49,7 @@ export const addModelsLoadedListener = (startAppListening: AppStartListening) =>
       handleLoRAModels(models, state, dispatch, log);
       handleControlAdapterModels(models, state, dispatch, log);
       handleSpandrelImageToImageModels(models, state, dispatch, log);
+      handleIPAdapterModels(models, state, dispatch, log);
     },
   });
 };
@@ -51,15 +58,15 @@ type ModelHandler = (
   models: AnyModelConfig[],
   state: RootState,
   dispatch: AppDispatch,
-  log: Logger<JSONObject>
+  log: Logger<SerializableObject>
 ) => undefined;
 
 const handleMainModels: ModelHandler = (models, state, dispatch, log) => {
-  const currentModel = state.generation.model;
+  const currentModel = state.params.model;
   const mainModels = models.filter(isNonRefinerMainModelConfig);
   if (mainModels.length === 0) {
     // No models loaded at all
-    dispatch(modelChanged(null));
+    dispatch(modelChanged({ model: null }));
     return;
   }
 
@@ -74,25 +81,16 @@ const handleMainModels: ModelHandler = (models, state, dispatch, log) => {
   if (defaultModelInList) {
     const result = zParameterModel.safeParse(defaultModelInList);
     if (result.success) {
-      dispatch(modelChanged(defaultModelInList, currentModel));
-
+      dispatch(modelChanged({ model: defaultModelInList, previousModel: currentModel }));
+      const { bbox } = selectCanvasSlice(state);
       const optimalDimension = getOptimalDimension(defaultModelInList);
-      if (
-        getIsSizeOptimal(
-          state.controlLayers.present.size.width,
-          state.controlLayers.present.size.height,
-          optimalDimension
-        )
-      ) {
+      if (getIsSizeOptimal(bbox.rect.width, bbox.rect.height, optimalDimension)) {
         return;
       }
-      const { width, height } = calculateNewSize(
-        state.controlLayers.present.size.aspectRatio.value,
-        optimalDimension * optimalDimension
-      );
+      const { width, height } = calculateNewSize(bbox.aspectRatio.value, optimalDimension * optimalDimension);
 
-      dispatch(widthChanged({ width }));
-      dispatch(heightChanged({ height }));
+      dispatch(bboxWidthChanged({ width }));
+      dispatch(bboxHeightChanged({ height }));
       return;
     }
   }
@@ -104,11 +102,11 @@ const handleMainModels: ModelHandler = (models, state, dispatch, log) => {
     return;
   }
 
-  dispatch(modelChanged(result.data, currentModel));
+  dispatch(modelChanged({ model: result.data, previousModel: currentModel }));
 };
 
 const handleRefinerModels: ModelHandler = (models, state, dispatch, _log) => {
-  const currentRefinerModel = state.sdxl.refinerModel;
+  const currentRefinerModel = state.params.refinerModel;
   const refinerModels = models.filter(isRefinerMainModelModelConfig);
   if (models.length === 0) {
     // No models loaded at all
@@ -127,7 +125,7 @@ const handleRefinerModels: ModelHandler = (models, state, dispatch, _log) => {
 };
 
 const handleVAEModels: ModelHandler = (models, state, dispatch, log) => {
-  const currentVae = state.generation.vae;
+  const currentVae = state.params.vae;
 
   if (currentVae === null) {
     // null is a valid VAE! it means "use the default with the main model"
@@ -160,28 +158,47 @@ const handleVAEModels: ModelHandler = (models, state, dispatch, log) => {
 };
 
 const handleLoRAModels: ModelHandler = (models, state, dispatch, _log) => {
-  const loras = state.lora.loras;
-
-  forEach(loras, (lora, id) => {
-    const isLoRAAvailable = models.some((m) => m.key === lora.model.key);
-
+  const loraModels = models.filter(isLoRAModelConfig);
+  state.loras.loras.forEach((lora) => {
+    const isLoRAAvailable = loraModels.some((m) => m.key === lora.model.key);
     if (isLoRAAvailable) {
       return;
     }
-
-    dispatch(loraRemoved(id));
+    dispatch(loraDeleted({ id: lora.id }));
   });
 };
 
 const handleControlAdapterModels: ModelHandler = (models, state, dispatch, _log) => {
-  selectControlAdapterAll(state.controlAdapters).forEach((ca) => {
-    const isModelAvailable = models.some((m) => m.key === ca.model?.key);
-
+  const caModels = models.filter(isControlNetOrT2IAdapterModelConfig);
+  selectCanvasSlice(state).controlLayers.entities.forEach((entity) => {
+    const isModelAvailable = caModels.some((m) => m.key === entity.controlAdapter.model?.key);
     if (isModelAvailable) {
       return;
     }
+    dispatch(controlLayerModelChanged({ entityIdentifier: getEntityIdentifier(entity), modelConfig: null }));
+  });
+};
 
-    dispatch(controlAdapterModelCleared({ id: ca.id }));
+const handleIPAdapterModels: ModelHandler = (models, state, dispatch, _log) => {
+  const ipaModels = models.filter(isIPAdapterModelConfig);
+  selectCanvasSlice(state).ipAdapters.entities.forEach((entity) => {
+    const isModelAvailable = ipaModels.some((m) => m.key === entity.ipAdapter.model?.key);
+    if (isModelAvailable) {
+      return;
+    }
+    dispatch(ipaModelChanged({ entityIdentifier: getEntityIdentifier(entity), modelConfig: null }));
+  });
+
+  selectCanvasSlice(state).regions.entities.forEach((entity) => {
+    entity.ipAdapters.forEach(({ id: ipAdapterId, model }) => {
+      const isModelAvailable = ipaModels.some((m) => m.key === model?.key);
+      if (isModelAvailable) {
+        return;
+      }
+      dispatch(
+        rgIPAdapterModelChanged({ entityIdentifier: getEntityIdentifier(entity), ipAdapterId, modelConfig: null })
+      );
+    });
   });
 };
 

invokeai/frontend/web/src/app/store/middleware/listenerMiddleware/listeners/promptChanged.ts

@@ -1,6 +1,6 @@
 import { isAnyOf } from '@reduxjs/toolkit';
 import type { AppStartListening } from 'app/store/middleware/listenerMiddleware';
-import { positivePromptChanged } from 'features/controlLayers/store/controlLayersSlice';
+import { positivePromptChanged } from 'features/controlLayers/store/paramsSlice';
 import {
   combinatorialToggled,
   isErrorChanged,
@@ -15,7 +15,7 @@ import { getPresetModifiedPrompts } from 'features/nodes/util/graph/graphBuilder
 import { activeStylePresetIdChanged } from 'features/stylePresets/store/stylePresetSlice';
 import { stylePresetsApi } from 'services/api/endpoints/stylePresets';
 import { utilitiesApi } from 'services/api/endpoints/utilities';
-import { socketConnected } from 'services/events/actions';
+import { socketConnected } from 'services/events/setEventListeners';
 
 const matcher = isAnyOf(
   positivePromptChanged,
@@ -24,8 +24,6 @@ const matcher = isAnyOf(
   maxPromptsReset,
   socketConnected,
   activeStylePresetIdChanged,
-  stylePresetsApi.endpoints.deleteStylePreset.matchFulfilled,
-  stylePresetsApi.endpoints.updateStylePreset.matchFulfilled,
   stylePresetsApi.endpoints.listStylePresets.matchFulfilled
 );
 

invokeai/frontend/web/src/app/store/middleware/listenerMiddleware/listeners/setDefaultSettings.ts

@@ -1,6 +1,5 @@
 import type { AppStartListening } from 'app/store/middleware/listenerMiddleware';
-import { heightChanged, widthChanged } from 'features/controlLayers/store/controlLayersSlice';
-import { setDefaultSettings } from 'features/parameters/store/actions';
+import { bboxHeightChanged, bboxWidthChanged } from 'features/controlLayers/store/canvasSlice';
 import {
   setCfgRescaleMultiplier,
   setCfgScale,
@@ -8,7 +7,8 @@ import {
   setSteps,
   vaePrecisionChanged,
   vaeSelected,
-} from 'features/parameters/store/generationSlice';
+} from 'features/controlLayers/store/paramsSlice';
+import { setDefaultSettings } from 'features/parameters/store/actions';
 import {
   isParameterCFGRescaleMultiplier,
   isParameterCFGScale,
@@ -30,7 +30,7 @@ export const addSetDefaultSettingsListener = (startAppListening: AppStartListeni
     effect: async (action, { dispatch, getState }) => {
       const state = getState();
 
-      const currentModel = state.generation.model;
+      const currentModel = state.params.model;
 
       if (!currentModel) {
         return;
@@ -98,13 +98,13 @@ export const addSetDefaultSettingsListener = (startAppListening: AppStartListeni
         const setSizeOptions = { updateAspectRatio: true, clamp: true };
         if (width) {
           if (isParameterWidth(width)) {
-            dispatch(widthChanged({ width, ...setSizeOptions }));
+            dispatch(bboxWidthChanged({ width, ...setSizeOptions }));
           }
         }
 
         if (height) {
           if (isParameterHeight(height)) {
-            dispatch(heightChanged({ height, ...setSizeOptions }));
+            dispatch(bboxHeightChanged({ height, ...setSizeOptions }));
           }
         }
 

invokeai/frontend/web/src/app/store/middleware/listenerMiddleware/listeners/socketConnected.ts

@@ -6,9 +6,9 @@ import { atom } from 'nanostores';
 import { api } from 'services/api';
 import { modelsApi } from 'services/api/endpoints/models';
 import { queueApi, selectQueueStatus } from 'services/api/endpoints/queue';
-import { socketConnected } from 'services/events/actions';
+import { socketConnected } from 'services/events/setEventListeners';
 
-const log = logger('socketio');
+const log = logger('events');
 
 const $isFirstConnection = atom(true);
 

invokeai/frontend/web/src/app/store/middleware/listenerMiddleware/listeners/socketio/socketDisconnected.ts

@@ -1,14 +0,0 @@
-import { logger } from 'app/logging/logger';
-import type { AppStartListening } from 'app/store/middleware/listenerMiddleware';
-import { socketDisconnected } from 'services/events/actions';
-
-const log = logger('socketio');
-
-export const addSocketDisconnectedEventListener = (startAppListening: AppStartListening) => {
-  startAppListening({
-    actionCreator: socketDisconnected,
-    effect: () => {
-      log.debug('Disconnected');
-    },
-  });
-};