chore: release v4.2.9.dev11

This provides some free reactivity for this canvas-manager-managed state.

.github/workflows/python-tests.yml

@@ -60,7 +60,7 @@ jobs:
             extra-index-url: 'https://download.pytorch.org/whl/cpu'
             github-env: $GITHUB_ENV
           - platform: macos-default
-            os: macOS-12
+            os: macOS-14
             github-env: $GITHUB_ENV
           - platform: windows-cpu
             os: windows-2022

docker/README.md

@@ -1,20 +1,22 @@
 # Invoke in Docker
 
-- Ensure that Docker can use the GPU on your system
-- This documentation assumes Linux, but should work similarly under Windows with WSL2
+First things first:
+
+- Ensure that Docker can use your [NVIDIA][nvidia docker docs] or [AMD][amd docker docs] GPU.
+- This document assumes a Linux system, but should work similarly under Windows with WSL2.
 - We don't recommend running Invoke in Docker on macOS at this time. It works, but very slowly.
 
-## Quickstart :lightning:
+## Quickstart
 
-No `docker compose`, no persistence, just a simple one-liner using the official images:
+No `docker compose`, no persistence, single command, using the official images:
 
-**CUDA:**
+**CUDA (NVIDIA GPU):**
 
 ```bash
 docker run --runtime=nvidia --gpus=all --publish 9090:9090 ghcr.io/invoke-ai/invokeai
 ```
 
-**ROCm:**
+**ROCm (AMD GPU):**
 
 ```bash
 docker run --device /dev/kfd --device /dev/dri --publish 9090:9090 ghcr.io/invoke-ai/invokeai:main-rocm
@@ -22,12 +24,20 @@ docker run --device /dev/kfd --device /dev/dri --publish 9090:9090 ghcr.io/invok
 
 Open `http://localhost:9090` in your browser once the container finishes booting, install some models, and generate away!
 
-> [!TIP]
-> To persist your data (including downloaded models) outside of the container, add a `--volume/-v` flag to the above command, e.g.: `docker run --volume /some/local/path:/invokeai <...the rest of the command>`
+### Data persistence
+
+To persist your generated images and downloaded models outside of the container, add a `--volume/-v` flag to the above command, e.g.:
+
+```bash
+docker run --volume /some/local/path:/invokeai {...etc...}
+```
+
+`/some/local/path/invokeai` will contain all your data.
+It can *usually* be reused between different installs of Invoke. Tread with caution and read the release notes!
 
 ## Customize the container
 
-We ship the `run.sh` script, which is a convenient wrapper around `docker compose` for cases where custom image build args are needed. Alternatively, the familiar `docker compose` commands work just as well.
+The included `run.sh` script is a convenience wrapper around `docker compose`. It can be helpful for passing additional build arguments to `docker compose`. Alternatively, the familiar `docker compose` commands work just as well.
 
 ```bash
 cd docker
@@ -38,11 +48,14 @@ cp .env.sample .env
 
 It will take a few minutes to build the image the first time. Once the application starts up, open `http://localhost:9090` in your browser to invoke!
 
+>[!TIP]
+>When using the `run.sh` script, the container will continue running after Ctrl+C. To shut it down, use the `docker compose down` command.
+
 ## Docker setup in detail
 
 #### Linux
 
-1. Ensure builkit is enabled in the Docker daemon settings (`/etc/docker/daemon.json`)
+1. Ensure buildkit is enabled in the Docker daemon settings (`/etc/docker/daemon.json`)
 2. Install the `docker compose` plugin using your package manager, or follow a [tutorial](https://docs.docker.com/compose/install/linux/#install-using-the-repository).
     - The deprecated `docker-compose` (hyphenated) CLI probably won't work. Update to a recent version.
 3. Ensure docker daemon is able to access the GPU.
@@ -98,25 +111,7 @@ GPU_DRIVER=cuda
 
 Any environment variables supported by InvokeAI can be set here. See the [Configuration docs](https://invoke-ai.github.io/InvokeAI/features/CONFIGURATION/) for further detail.
 
-## Even More Customizing!
+---
 
-See the `docker-compose.yml` file. The `command` instruction can be uncommented and used to run arbitrary startup commands. Some examples below.
-
-### Reconfigure the runtime directory
-
-Can be used to download additional models from the supported model list
-
-In conjunction with `INVOKEAI_ROOT` can be also used to initialize a runtime directory
-
-```yaml
-command:
-  - invokeai-configure
-  - --yes
-```
-
-Or install models:
-
-```yaml
-command:
-  - invokeai-model-install
-```
+[nvidia docker docs]: https://docs.nvidia.com/datacenter/cloud-native/container-toolkit/latest/install-guide.html
+[amd docker docs]: https://rocm.docs.amd.com/projects/install-on-linux/en/latest/how-to/docker.html

invokeai/app/api/routers/session_queue.py

@@ -11,6 +11,7 @@ from invokeai.app.services.session_queue.session_queue_common import (
     Batch,
     BatchStatus,
     CancelByBatchIDsResult,
+    CancelByOriginResult,
     ClearResult,
     EnqueueBatchResult,
     PruneResult,
@@ -105,6 +106,19 @@ 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_origin",
+    operation_id="cancel_by_origin",
+    responses={200: {"model": CancelByBatchIDsResult}},
+)
+async def cancel_by_origin(
+    queue_id: str = Path(description="The queue id to perform this operation on"),
+    origin: str = Query(description="The origin to cancel all queue items for"),
+) -> CancelByOriginResult:
+    """Immediately cancels all queue items with the given origin"""
+    return ApiDependencies.invoker.services.session_queue.cancel_by_origin(queue_id=queue_id, origin=origin)
+
+
 @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/fields.py

@@ -40,14 +40,18 @@ class UIType(str, Enum, metaclass=MetaEnum):
 
     # region Model Field Types
     MainModel = "MainModelField"
+    FluxMainModel = "FluxMainModelField"
     SDXLMainModel = "SDXLMainModelField"
     SDXLRefinerModel = "SDXLRefinerModelField"
     ONNXModel = "ONNXModelField"
     VAEModel = "VAEModelField"
+    FluxVAEModel = "FluxVAEModelField"
     LoRAModel = "LoRAModelField"
     ControlNetModel = "ControlNetModelField"
     IPAdapterModel = "IPAdapterModelField"
     T2IAdapterModel = "T2IAdapterModelField"
+    T5EncoderModel = "T5EncoderModelField"
+    CLIPEmbedModel = "CLIPEmbedModelField"
     SpandrelImageToImageModel = "SpandrelImageToImageModelField"
     # endregion
 
@@ -125,13 +129,17 @@ class FieldDescriptions:
     negative_cond = "Negative conditioning tensor"
     noise = "Noise tensor"
     clip = "CLIP (tokenizer, text encoder, LoRAs) and skipped layer count"
+    t5_encoder = "T5 tokenizer and text encoder"
+    clip_embed_model = "CLIP Embed loader"
     unet = "UNet (scheduler, LoRAs)"
+    transformer = "Transformer"
     vae = "VAE"
     cond = "Conditioning tensor"
     controlnet_model = "ControlNet model to load"
     vae_model = "VAE model to load"
     lora_model = "LoRA model to load"
     main_model = "Main model (UNet, VAE, CLIP) to load"
+    flux_model = "Flux model (Transformer) to load"
     sdxl_main_model = "SDXL Main model (UNet, VAE, CLIP1, CLIP2) to load"
     sdxl_refiner_model = "SDXL Refiner Main Modde (UNet, VAE, CLIP2) to load"
     onnx_main_model = "ONNX Main model (UNet, VAE, CLIP) to load"
@@ -231,6 +239,12 @@ class ColorField(BaseModel):
         return (self.r, self.g, self.b, self.a)
 
 
+class FluxConditioningField(BaseModel):
+    """A conditioning tensor primitive value"""
+
+    conditioning_name: str = Field(description="The name of conditioning tensor")
+
+
 class ConditioningField(BaseModel):
     """A conditioning tensor primitive value"""
 

invokeai/app/invocations/flux_text_encoder.py

@@ -0,0 +1,92 @@
+from typing import Literal
+
+import torch
+from transformers import CLIPTextModel, CLIPTokenizer, T5EncoderModel, T5Tokenizer
+
+from invokeai.app.invocations.baseinvocation import BaseInvocation, Classification, invocation
+from invokeai.app.invocations.fields import FieldDescriptions, Input, InputField
+from invokeai.app.invocations.model import CLIPField, T5EncoderField
+from invokeai.app.invocations.primitives import FluxConditioningOutput
+from invokeai.app.services.shared.invocation_context import InvocationContext
+from invokeai.backend.flux.modules.conditioner import HFEncoder
+from invokeai.backend.stable_diffusion.diffusion.conditioning_data import ConditioningFieldData, FLUXConditioningInfo
+
+
+@invocation(
+    "flux_text_encoder",
+    title="FLUX Text Encoding",
+    tags=["prompt", "conditioning", "flux"],
+    category="conditioning",
+    version="1.0.0",
+    classification=Classification.Prototype,
+)
+class FluxTextEncoderInvocation(BaseInvocation):
+    """Encodes and preps a prompt for a flux image."""
+
+    clip: CLIPField = InputField(
+        title="CLIP",
+        description=FieldDescriptions.clip,
+        input=Input.Connection,
+    )
+    t5_encoder: T5EncoderField = InputField(
+        title="T5Encoder",
+        description=FieldDescriptions.t5_encoder,
+        input=Input.Connection,
+    )
+    t5_max_seq_len: Literal[256, 512] = InputField(
+        description="Max sequence length for the T5 encoder. Expected to be 256 for FLUX schnell models and 512 for FLUX dev models."
+    )
+    prompt: str = InputField(description="Text prompt to encode.")
+
+    @torch.no_grad()
+    def invoke(self, context: InvocationContext) -> FluxConditioningOutput:
+        # Note: The T5 and CLIP encoding are done in separate functions to ensure that all model references are locally
+        # scoped. This ensures that the T5 model can be freed and gc'd before loading the CLIP model (if necessary).
+        t5_embeddings = self._t5_encode(context)
+        clip_embeddings = self._clip_encode(context)
+        conditioning_data = ConditioningFieldData(
+            conditionings=[FLUXConditioningInfo(clip_embeds=clip_embeddings, t5_embeds=t5_embeddings)]
+        )
+
+        conditioning_name = context.conditioning.save(conditioning_data)
+        return FluxConditioningOutput.build(conditioning_name)
+
+    def _t5_encode(self, context: InvocationContext) -> torch.Tensor:
+        t5_tokenizer_info = context.models.load(self.t5_encoder.tokenizer)
+        t5_text_encoder_info = context.models.load(self.t5_encoder.text_encoder)
+
+        prompt = [self.prompt]
+
+        with (
+            t5_text_encoder_info as t5_text_encoder,
+            t5_tokenizer_info as t5_tokenizer,
+        ):
+            assert isinstance(t5_text_encoder, T5EncoderModel)
+            assert isinstance(t5_tokenizer, T5Tokenizer)
+
+            t5_encoder = HFEncoder(t5_text_encoder, t5_tokenizer, False, self.t5_max_seq_len)
+
+            prompt_embeds = t5_encoder(prompt)
+
+        assert isinstance(prompt_embeds, torch.Tensor)
+        return prompt_embeds
+
+    def _clip_encode(self, context: InvocationContext) -> torch.Tensor:
+        clip_tokenizer_info = context.models.load(self.clip.tokenizer)
+        clip_text_encoder_info = context.models.load(self.clip.text_encoder)
+
+        prompt = [self.prompt]
+
+        with (
+            clip_text_encoder_info as clip_text_encoder,
+            clip_tokenizer_info as clip_tokenizer,
+        ):
+            assert isinstance(clip_text_encoder, CLIPTextModel)
+            assert isinstance(clip_tokenizer, CLIPTokenizer)
+
+            clip_encoder = HFEncoder(clip_text_encoder, clip_tokenizer, True, 77)
+
+            pooled_prompt_embeds = clip_encoder(prompt)
+
+        assert isinstance(pooled_prompt_embeds, torch.Tensor)
+        return pooled_prompt_embeds

invokeai/app/invocations/flux_text_to_image.py

@@ -0,0 +1,169 @@
+import torch
+from einops import rearrange
+from PIL import Image
+
+from invokeai.app.invocations.baseinvocation import BaseInvocation, Classification, invocation
+from invokeai.app.invocations.fields import (
+    FieldDescriptions,
+    FluxConditioningField,
+    Input,
+    InputField,
+    WithBoard,
+    WithMetadata,
+)
+from invokeai.app.invocations.model import TransformerField, VAEField
+from invokeai.app.invocations.primitives import ImageOutput
+from invokeai.app.services.session_processor.session_processor_common import CanceledException
+from invokeai.app.services.shared.invocation_context import InvocationContext
+from invokeai.backend.flux.model import Flux
+from invokeai.backend.flux.modules.autoencoder import AutoEncoder
+from invokeai.backend.flux.sampling import denoise, get_noise, get_schedule, prepare_latent_img_patches, unpack
+from invokeai.backend.stable_diffusion.diffusion.conditioning_data import FLUXConditioningInfo
+from invokeai.backend.util.devices import TorchDevice
+
+
+@invocation(
+    "flux_text_to_image",
+    title="FLUX Text to Image",
+    tags=["image", "flux"],
+    category="image",
+    version="1.0.0",
+    classification=Classification.Prototype,
+)
+class FluxTextToImageInvocation(BaseInvocation, WithMetadata, WithBoard):
+    """Text-to-image generation using a FLUX model."""
+
+    transformer: TransformerField = InputField(
+        description=FieldDescriptions.flux_model,
+        input=Input.Connection,
+        title="Transformer",
+    )
+    vae: VAEField = InputField(
+        description=FieldDescriptions.vae,
+        input=Input.Connection,
+    )
+    positive_text_conditioning: FluxConditioningField = InputField(
+        description=FieldDescriptions.positive_cond, input=Input.Connection
+    )
+    width: int = InputField(default=1024, multiple_of=16, description="Width of the generated image.")
+    height: int = InputField(default=1024, multiple_of=16, description="Height of the generated image.")
+    num_steps: int = InputField(
+        default=4, description="Number of diffusion steps. Recommend values are schnell: 4, dev: 50."
+    )
+    guidance: float = InputField(
+        default=4.0,
+        description="The guidance strength. Higher values adhere more strictly to the prompt, and will produce less diverse images. FLUX dev only, ignored for schnell.",
+    )
+    seed: int = InputField(default=0, description="Randomness seed for reproducibility.")
+
+    @torch.no_grad()
+    def invoke(self, context: InvocationContext) -> ImageOutput:
+        latents = self._run_diffusion(context)
+        image = self._run_vae_decoding(context, latents)
+        image_dto = context.images.save(image=image)
+        return ImageOutput.build(image_dto)
+
+    def _run_diffusion(
+        self,
+        context: InvocationContext,
+    ):
+        inference_dtype = torch.bfloat16
+
+        # Load the conditioning data.
+        cond_data = context.conditioning.load(self.positive_text_conditioning.conditioning_name)
+        assert len(cond_data.conditionings) == 1
+        flux_conditioning = cond_data.conditionings[0]
+        assert isinstance(flux_conditioning, FLUXConditioningInfo)
+        flux_conditioning = flux_conditioning.to(dtype=inference_dtype)
+        t5_embeddings = flux_conditioning.t5_embeds
+        clip_embeddings = flux_conditioning.clip_embeds
+
+        transformer_info = context.models.load(self.transformer.transformer)
+
+        # Prepare input noise.
+        x = get_noise(
+            num_samples=1,
+            height=self.height,
+            width=self.width,
+            device=TorchDevice.choose_torch_device(),
+            dtype=inference_dtype,
+            seed=self.seed,
+        )
+
+        x, img_ids = prepare_latent_img_patches(x)
+
+        is_schnell = "schnell" in transformer_info.config.config_path
+
+        timesteps = get_schedule(
+            num_steps=self.num_steps,
+            image_seq_len=x.shape[1],
+            shift=not is_schnell,
+        )
+
+        bs, t5_seq_len, _ = t5_embeddings.shape
+        txt_ids = torch.zeros(bs, t5_seq_len, 3, dtype=inference_dtype, device=TorchDevice.choose_torch_device())
+
+        with transformer_info as transformer:
+            assert isinstance(transformer, Flux)
+
+            def step_callback() -> None:
+                if context.util.is_canceled():
+                    raise CanceledException
+
+                # TODO: Make this look like the image before re-enabling
+                # latent_image = unpack(img.float(), self.height, self.width)
+                # latent_image = latent_image.squeeze()  # Remove unnecessary dimensions
+                # flattened_tensor = latent_image.reshape(-1)  # Flatten to shape [48*128*128]
+
+                # # Create a new tensor of the required shape [255, 255, 3]
+                # latent_image = flattened_tensor[: 255 * 255 * 3].reshape(255, 255, 3)  # Reshape to RGB format
+
+                # # Convert to a NumPy array and then to a PIL Image
+                # image = Image.fromarray(latent_image.cpu().numpy().astype(np.uint8))
+
+                # (width, height) = image.size
+                # width *= 8
+                # height *= 8
+
+                # dataURL = image_to_dataURL(image, image_format="JPEG")
+
+                # # TODO: move this whole function to invocation context to properly reference these variables
+                # context._services.events.emit_invocation_denoise_progress(
+                #     context._data.queue_item,
+                #     context._data.invocation,
+                #     state,
+                #     ProgressImage(dataURL=dataURL, width=width, height=height),
+                # )
+
+            x = denoise(
+                model=transformer,
+                img=x,
+                img_ids=img_ids,
+                txt=t5_embeddings,
+                txt_ids=txt_ids,
+                vec=clip_embeddings,
+                timesteps=timesteps,
+                step_callback=step_callback,
+                guidance=self.guidance,
+            )
+
+        x = unpack(x.float(), self.height, self.width)
+
+        return x
+
+    def _run_vae_decoding(
+        self,
+        context: InvocationContext,
+        latents: torch.Tensor,
+    ) -> Image.Image:
+        vae_info = context.models.load(self.vae.vae)
+        with vae_info as vae:
+            assert isinstance(vae, AutoEncoder)
+            latents = latents.to(dtype=TorchDevice.choose_torch_dtype())
+            img = vae.decode(latents)
+
+        img = img.clamp(-1, 1)
+        img = rearrange(img[0], "c h w -> h w c")
+        img_pil = Image.fromarray((127.5 * (img + 1.0)).byte().cpu().numpy())
+
+        return img_pil

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

@@ -1,5 +1,5 @@
 import copy
-from typing import List, Optional
+from typing import List, Literal, Optional
 
 from pydantic import BaseModel, Field
 
@@ -13,7 +13,14 @@ from invokeai.app.invocations.baseinvocation import (
 from invokeai.app.invocations.fields import FieldDescriptions, Input, InputField, OutputField, UIType
 from invokeai.app.services.shared.invocation_context import InvocationContext
 from invokeai.app.shared.models import FreeUConfig
-from invokeai.backend.model_manager.config import AnyModelConfig, BaseModelType, ModelType, SubModelType
+from invokeai.backend.flux.util import max_seq_lengths
+from invokeai.backend.model_manager.config import (
+    AnyModelConfig,
+    BaseModelType,
+    CheckpointConfigBase,
+    ModelType,
+    SubModelType,
+)
 
 
 class ModelIdentifierField(BaseModel):
@@ -60,6 +67,15 @@ class CLIPField(BaseModel):
     loras: List[LoRAField] = Field(description="LoRAs to apply on model loading")
 
 
+class TransformerField(BaseModel):
+    transformer: ModelIdentifierField = Field(description="Info to load Transformer submodel")
+
+
+class T5EncoderField(BaseModel):
+    tokenizer: ModelIdentifierField = Field(description="Info to load tokenizer submodel")
+    text_encoder: ModelIdentifierField = Field(description="Info to load text_encoder submodel")
+
+
 class VAEField(BaseModel):
     vae: ModelIdentifierField = Field(description="Info to load vae submodel")
     seamless_axes: List[str] = Field(default_factory=list, description='Axes("x" and "y") to which apply seamless')
@@ -122,6 +138,78 @@ class ModelIdentifierInvocation(BaseInvocation):
         return ModelIdentifierOutput(model=self.model)
 
 
+@invocation_output("flux_model_loader_output")
+class FluxModelLoaderOutput(BaseInvocationOutput):
+    """Flux base model loader output"""
+
+    transformer: TransformerField = OutputField(description=FieldDescriptions.transformer, title="Transformer")
+    clip: CLIPField = OutputField(description=FieldDescriptions.clip, title="CLIP")
+    t5_encoder: T5EncoderField = OutputField(description=FieldDescriptions.t5_encoder, title="T5 Encoder")
+    vae: VAEField = OutputField(description=FieldDescriptions.vae, title="VAE")
+    max_seq_len: Literal[256, 512] = OutputField(
+        description="The max sequence length to used for the T5 encoder. (256 for schnell transformer, 512 for dev transformer)",
+        title="Max Seq Length",
+    )
+
+
+@invocation(
+    "flux_model_loader",
+    title="Flux Main Model",
+    tags=["model", "flux"],
+    category="model",
+    version="1.0.4",
+    classification=Classification.Prototype,
+)
+class FluxModelLoaderInvocation(BaseInvocation):
+    """Loads a flux base model, outputting its submodels."""
+
+    model: ModelIdentifierField = InputField(
+        description=FieldDescriptions.flux_model,
+        ui_type=UIType.FluxMainModel,
+        input=Input.Direct,
+    )
+
+    t5_encoder_model: ModelIdentifierField = InputField(
+        description=FieldDescriptions.t5_encoder, ui_type=UIType.T5EncoderModel, input=Input.Direct, title="T5 Encoder"
+    )
+
+    clip_embed_model: ModelIdentifierField = InputField(
+        description=FieldDescriptions.clip_embed_model,
+        ui_type=UIType.CLIPEmbedModel,
+        input=Input.Direct,
+        title="CLIP Embed",
+    )
+
+    vae_model: ModelIdentifierField = InputField(
+        description=FieldDescriptions.vae_model, ui_type=UIType.FluxVAEModel, title="VAE"
+    )
+
+    def invoke(self, context: InvocationContext) -> FluxModelLoaderOutput:
+        for key in [self.model.key, self.t5_encoder_model.key, self.clip_embed_model.key, self.vae_model.key]:
+            if not context.models.exists(key):
+                raise ValueError(f"Unknown model: {key}")
+
+        transformer = self.model.model_copy(update={"submodel_type": SubModelType.Transformer})
+        vae = self.vae_model.model_copy(update={"submodel_type": SubModelType.VAE})
+
+        tokenizer = self.clip_embed_model.model_copy(update={"submodel_type": SubModelType.Tokenizer})
+        clip_encoder = self.clip_embed_model.model_copy(update={"submodel_type": SubModelType.TextEncoder})
+
+        tokenizer2 = self.t5_encoder_model.model_copy(update={"submodel_type": SubModelType.Tokenizer2})
+        t5_encoder = self.t5_encoder_model.model_copy(update={"submodel_type": SubModelType.TextEncoder2})
+
+        transformer_config = context.models.get_config(transformer)
+        assert isinstance(transformer_config, CheckpointConfigBase)
+
+        return FluxModelLoaderOutput(
+            transformer=TransformerField(transformer=transformer),
+            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),
+            max_seq_len=max_seq_lengths[transformer_config.config_path],
+        )
+
+
 @invocation(
     "main_model_loader",
     title="Main Model",

invokeai/app/invocations/primitives.py

@@ -12,6 +12,7 @@ from invokeai.app.invocations.fields import (
     ConditioningField,
     DenoiseMaskField,
     FieldDescriptions,
+    FluxConditioningField,
     ImageField,
     Input,
     InputField,
@@ -414,6 +415,17 @@ class MaskOutput(BaseInvocationOutput):
     height: int = OutputField(description="The height of the mask in pixels.")
 
 
+@invocation_output("flux_conditioning_output")
+class FluxConditioningOutput(BaseInvocationOutput):
+    """Base class for nodes that output a single conditioning tensor"""
+
+    conditioning: FluxConditioningField = OutputField(description=FieldDescriptions.cond)
+
+    @classmethod
+    def build(cls, conditioning_name: str) -> "FluxConditioningOutput":
+        return cls(conditioning=FluxConditioningField(conditioning_name=conditioning_name))
+
+
 @invocation_output("conditioning_output")
 class ConditioningOutput(BaseInvocationOutput):
     """Base class for nodes that output a single conditioning tensor"""

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/model_install/model_install_default.py

@@ -783,8 +783,9 @@ class ModelInstallService(ModelInstallServiceBase):
         # So what we do is to synthesize a folder named "sdxl-turbo_vae" here.
         if subfolder:
             top = Path(remote_files[0].path.parts[0])  # e.g. "sdxl-turbo/"
-            path_to_remove = top / subfolder.parts[-1]  # sdxl-turbo/vae/
-            path_to_add = Path(f"{top}_{subfolder}")
+            path_to_remove = top / subfolder  # sdxl-turbo/vae/
+            subfolder_rename = subfolder.name.replace("/", "_").replace("\\", "_")
+            path_to_add = Path(f"{top}_{subfolder_rename}")
         else:
             path_to_remove = Path(".")
             path_to_add = Path(".")

invokeai/app/services/model_records/model_records_base.py

@@ -77,6 +77,7 @@ class ModelRecordChanges(BaseModelExcludeNull):
     type: Optional[ModelType] = Field(description="Type of model", default=None)
     key: Optional[str] = Field(description="Database ID for this model", default=None)
     hash: Optional[str] = Field(description="hash of model file", default=None)
+    format: Optional[str] = Field(description="format of model file", default=None)
     trigger_phrases: Optional[set[str]] = Field(description="Set of trigger phrases for this model", default=None)
     default_settings: Optional[MainModelDefaultSettings | ControlAdapterDefaultSettings] = Field(
         description="Default settings for this model", default=None

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,
+    CancelByOriginResult,
     CancelByQueueIDResult,
     ClearResult,
     EnqueueBatchResult,
@@ -95,6 +96,11 @@ class SessionQueueBase(ABC):
         """Cancels all queue items with matching batch IDs"""
         pass
 
+    @abstractmethod
+    def cancel_by_origin(self, queue_id: str, origin: str) -> CancelByOriginResult:
+        """Cancels all queue items with the given batch origin"""
+        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 CancelByOriginResult(BaseModel):
+    """Result of canceling by list of batch ids"""
+
+    canceled: int = Field(..., description="Number of queue items canceled")
+
+
 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,
+    CancelByOriginResult,
     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_origin(self, queue_id: str, origin: str) -> CancelByOriginResult:
+        try:
+            current_queue_item = self.get_current(queue_id)
+            self.__lock.acquire()
+            where = """--sql
+                WHERE
+                  queue_id == ?
+                  AND origin == ?
+                  AND status != 'canceled'
+                  AND status != 'completed'
+                  AND status != 'failed'
+                """
+            params = (queue_id, origin)
+            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.origin == origin:
+                self._set_queue_item_status(current_queue_item.item_id, "canceled")
+        except Exception:
+            self.__conn.rollback()
+            raise
+        finally:
+            self.__lock.release()
+        return CancelByOriginResult(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/app/services/workflow_records/default_workflows/Flux Text to Image.json

@@ -0,0 +1,260 @@
+{
+  "name": "FLUX Text to Image",
+  "author": "InvokeAI",
+  "description": "A simple text-to-image workflow using FLUX dev or schnell models. Prerequisite model downloads: T5 Encoder, CLIP-L Encoder, and FLUX VAE. Quantized and un-quantized versions can be found in the starter models tab within your Model Manager. We recommend 4 steps for FLUX schnell models and 30 steps for FLUX dev models.",
+  "version": "1.0.4",
+  "contact": "",
+  "tags": "text2image, flux",
+  "notes": "Prerequisite model downloads: T5 Encoder, CLIP-L Encoder, and FLUX VAE. Quantized and un-quantized versions can be found in the starter models tab within your Model Manager. We recommend 4 steps for FLUX schnell models and 30 steps for FLUX dev models.",
+  "exposedFields": [
+    {
+      "nodeId": "f8d9d7c8-9ed7-4bd7-9e42-ab0e89bfac90",
+      "fieldName": "model"
+    },
+    {
+      "nodeId": "01f674f8-b3d1-4df1-acac-6cb8e0bfb63c",
+      "fieldName": "prompt"
+    },
+    {
+      "nodeId": "159bdf1b-79e7-4174-b86e-d40e646964c8",
+      "fieldName": "num_steps"
+    },
+    {
+      "nodeId": "f8d9d7c8-9ed7-4bd7-9e42-ab0e89bfac90",
+      "fieldName": "t5_encoder_model"
+    }
+  ],
+  "meta": {
+    "version": "3.0.0",
+    "category": "default"
+  },
+  "nodes": [
+    {
+      "id": "f8d9d7c8-9ed7-4bd7-9e42-ab0e89bfac90",
+      "type": "invocation",
+      "data": {
+        "id": "f8d9d7c8-9ed7-4bd7-9e42-ab0e89bfac90",
+        "type": "flux_model_loader",
+        "version": "1.0.4",
+        "label": "",
+        "notes": "",
+        "isOpen": true,
+        "isIntermediate": true,
+        "useCache": false,
+        "inputs": {
+          "model": {
+            "name": "model",
+            "label": ""
+          },
+          "t5_encoder_model": {
+            "name": "t5_encoder_model",
+            "label": ""
+          },
+          "clip_embed_model": {
+            "name": "clip_embed_model",
+            "label": ""
+          },
+          "vae_model": {
+            "name": "vae_model",
+            "label": ""
+          }
+        }
+      },
+      "position": {
+        "x": 381.1882713063478,
+        "y": -95.89663532854017
+      }
+    },
+    {
+      "id": "01f674f8-b3d1-4df1-acac-6cb8e0bfb63c",
+      "type": "invocation",
+      "data": {
+        "id": "01f674f8-b3d1-4df1-acac-6cb8e0bfb63c",
+        "type": "flux_text_encoder",
+        "version": "1.0.0",
+        "label": "",
+        "notes": "",
+        "isOpen": true,
+        "isIntermediate": true,
+        "useCache": true,
+        "inputs": {
+          "clip": {
+            "name": "clip",
+            "label": ""
+          },
+          "t5_encoder": {
+            "name": "t5_encoder",
+            "label": ""
+          },
+          "t5_max_seq_len": {
+            "name": "t5_max_seq_len",
+            "label": "T5 Max Seq Len",
+            "value": 256
+          },
+          "prompt": {
+            "name": "prompt",
+            "label": "",
+            "value": "a cat"
+          }
+        }
+      },
+      "position": {
+        "x": 824.1970602278849,
+        "y": 146.98251001061735
+      }
+    },
+    {
+      "id": "4754c534-a5f3-4ad0-9382-7887985e668c",
+      "type": "invocation",
+      "data": {
+        "id": "4754c534-a5f3-4ad0-9382-7887985e668c",
+        "type": "rand_int",
+        "version": "1.0.1",
+        "label": "",
+        "notes": "",
+        "isOpen": true,
+        "isIntermediate": true,
+        "useCache": false,
+        "inputs": {
+          "low": {
+            "name": "low",
+            "label": "",
+            "value": 0
+          },
+          "high": {
+            "name": "high",
+            "label": "",
+            "value": 2147483647
+          }
+        }
+      },
+      "position": {
+        "x": 822.9899179655476,
+        "y": 360.9657214885052
+      }
+    },
+    {
+      "id": "159bdf1b-79e7-4174-b86e-d40e646964c8",
+      "type": "invocation",
+      "data": {
+        "id": "159bdf1b-79e7-4174-b86e-d40e646964c8",
+        "type": "flux_text_to_image",
+        "version": "1.0.0",
+        "label": "",
+        "notes": "",
+        "isOpen": true,
+        "isIntermediate": false,
+        "useCache": true,
+        "inputs": {
+          "board": {
+            "name": "board",
+            "label": ""
+          },
+          "metadata": {
+            "name": "metadata",
+            "label": ""
+          },
+          "transformer": {
+            "name": "transformer",
+            "label": ""
+          },
+          "vae": {
+            "name": "vae",
+            "label": ""
+          },
+          "positive_text_conditioning": {
+            "name": "positive_text_conditioning",
+            "label": ""
+          },
+          "width": {
+            "name": "width",
+            "label": "",
+            "value": 1024
+          },
+          "height": {
+            "name": "height",
+            "label": "",
+            "value": 1024
+          },
+          "num_steps": {
+            "name": "num_steps",
+            "label": "Steps (Recommend 30 for Dev, 4 for Schnell)",
+            "value": 30
+          },
+          "guidance": {
+            "name": "guidance",
+            "label": "",
+            "value": 4
+          },
+          "seed": {
+            "name": "seed",
+            "label": "",
+            "value": 0
+          }
+        }
+      },
+      "position": {
+        "x": 1216.3900791301849,
+        "y": 5.500841807102248
+      }
+    }
+  ],
+  "edges": [
+    {
+      "id": "reactflow__edge-f8d9d7c8-9ed7-4bd7-9e42-ab0e89bfac90max_seq_len-01f674f8-b3d1-4df1-acac-6cb8e0bfb63ct5_max_seq_len",
+      "type": "default",
+      "source": "f8d9d7c8-9ed7-4bd7-9e42-ab0e89bfac90",
+      "target": "01f674f8-b3d1-4df1-acac-6cb8e0bfb63c",
+      "sourceHandle": "max_seq_len",
+      "targetHandle": "t5_max_seq_len"
+    },
+    {
+      "id": "reactflow__edge-f8d9d7c8-9ed7-4bd7-9e42-ab0e89bfac90vae-159bdf1b-79e7-4174-b86e-d40e646964c8vae",
+      "type": "default",
+      "source": "f8d9d7c8-9ed7-4bd7-9e42-ab0e89bfac90",
+      "target": "159bdf1b-79e7-4174-b86e-d40e646964c8",
+      "sourceHandle": "vae",
+      "targetHandle": "vae"
+    },
+    {
+      "id": "reactflow__edge-f8d9d7c8-9ed7-4bd7-9e42-ab0e89bfac90t5_encoder-01f674f8-b3d1-4df1-acac-6cb8e0bfb63ct5_encoder",
+      "type": "default",
+      "source": "f8d9d7c8-9ed7-4bd7-9e42-ab0e89bfac90",
+      "target": "01f674f8-b3d1-4df1-acac-6cb8e0bfb63c",
+      "sourceHandle": "t5_encoder",
+      "targetHandle": "t5_encoder"
+    },
+    {
+      "id": "reactflow__edge-f8d9d7c8-9ed7-4bd7-9e42-ab0e89bfac90clip-01f674f8-b3d1-4df1-acac-6cb8e0bfb63cclip",
+      "type": "default",
+      "source": "f8d9d7c8-9ed7-4bd7-9e42-ab0e89bfac90",
+      "target": "01f674f8-b3d1-4df1-acac-6cb8e0bfb63c",
+      "sourceHandle": "clip",
+      "targetHandle": "clip"
+    },
+    {
+      "id": "reactflow__edge-f8d9d7c8-9ed7-4bd7-9e42-ab0e89bfac90transformer-159bdf1b-79e7-4174-b86e-d40e646964c8transformer",
+      "type": "default",
+      "source": "f8d9d7c8-9ed7-4bd7-9e42-ab0e89bfac90",
+      "target": "159bdf1b-79e7-4174-b86e-d40e646964c8",
+      "sourceHandle": "transformer",
+      "targetHandle": "transformer"
+    },
+    {
+      "id": "reactflow__edge-01f674f8-b3d1-4df1-acac-6cb8e0bfb63cconditioning-159bdf1b-79e7-4174-b86e-d40e646964c8positive_text_conditioning",
+      "type": "default",
+      "source": "01f674f8-b3d1-4df1-acac-6cb8e0bfb63c",
+      "target": "159bdf1b-79e7-4174-b86e-d40e646964c8",
+      "sourceHandle": "conditioning",
+      "targetHandle": "positive_text_conditioning"
+    },
+    {
+      "id": "reactflow__edge-4754c534-a5f3-4ad0-9382-7887985e668cvalue-159bdf1b-79e7-4174-b86e-d40e646964c8seed",
+      "type": "default",
+      "source": "4754c534-a5f3-4ad0-9382-7887985e668c",
+      "target": "159bdf1b-79e7-4174-b86e-d40e646964c8",
+      "sourceHandle": "value",
+      "targetHandle": "seed"
+    }
+  ]
+}

invokeai/backend/flux/math.py

@@ -0,0 +1,32 @@
+# Initially pulled from https://github.com/black-forest-labs/flux
+
+import torch
+from einops import rearrange
+from torch import Tensor
+
+
+def attention(q: Tensor, k: Tensor, v: Tensor, pe: Tensor) -> Tensor:
+    q, k = apply_rope(q, k, pe)
+
+    x = torch.nn.functional.scaled_dot_product_attention(q, k, v)
+    x = rearrange(x, "B H L D -> B L (H D)")
+
+    return x
+
+
+def rope(pos: Tensor, dim: int, theta: int) -> Tensor:
+    assert dim % 2 == 0
+    scale = torch.arange(0, dim, 2, dtype=torch.float64, device=pos.device) / dim
+    omega = 1.0 / (theta**scale)
+    out = torch.einsum("...n,d->...nd", pos, omega)
+    out = torch.stack([torch.cos(out), -torch.sin(out), torch.sin(out), torch.cos(out)], dim=-1)
+    out = rearrange(out, "b n d (i j) -> b n d i j", i=2, j=2)
+    return out.float()
+
+
+def apply_rope(xq: Tensor, xk: Tensor, freqs_cis: Tensor) -> tuple[Tensor, Tensor]:
+    xq_ = xq.float().reshape(*xq.shape[:-1], -1, 1, 2)
+    xk_ = xk.float().reshape(*xk.shape[:-1], -1, 1, 2)
+    xq_out = freqs_cis[..., 0] * xq_[..., 0] + freqs_cis[..., 1] * xq_[..., 1]
+    xk_out = freqs_cis[..., 0] * xk_[..., 0] + freqs_cis[..., 1] * xk_[..., 1]
+    return xq_out.reshape(*xq.shape).type_as(xq), xk_out.reshape(*xk.shape).type_as(xk)

invokeai/backend/flux/model.py

@@ -0,0 +1,117 @@
+# Initially pulled from https://github.com/black-forest-labs/flux
+
+from dataclasses import dataclass
+
+import torch
+from torch import Tensor, nn
+
+from invokeai.backend.flux.modules.layers import (
+    DoubleStreamBlock,
+    EmbedND,
+    LastLayer,
+    MLPEmbedder,
+    SingleStreamBlock,
+    timestep_embedding,
+)
+
+
+@dataclass
+class FluxParams:
+    in_channels: int
+    vec_in_dim: int
+    context_in_dim: int
+    hidden_size: int
+    mlp_ratio: float
+    num_heads: int
+    depth: int
+    depth_single_blocks: int
+    axes_dim: list[int]
+    theta: int
+    qkv_bias: bool
+    guidance_embed: bool
+
+
+class Flux(nn.Module):
+    """
+    Transformer model for flow matching on sequences.
+    """
+
+    def __init__(self, params: FluxParams):
+        super().__init__()
+
+        self.params = params
+        self.in_channels = params.in_channels
+        self.out_channels = self.in_channels
+        if params.hidden_size % params.num_heads != 0:
+            raise ValueError(f"Hidden size {params.hidden_size} must be divisible by num_heads {params.num_heads}")
+        pe_dim = params.hidden_size // params.num_heads
+        if sum(params.axes_dim) != pe_dim:
+            raise ValueError(f"Got {params.axes_dim} but expected positional dim {pe_dim}")
+        self.hidden_size = params.hidden_size
+        self.num_heads = params.num_heads
+        self.pe_embedder = EmbedND(dim=pe_dim, theta=params.theta, axes_dim=params.axes_dim)
+        self.img_in = nn.Linear(self.in_channels, self.hidden_size, bias=True)
+        self.time_in = MLPEmbedder(in_dim=256, hidden_dim=self.hidden_size)
+        self.vector_in = MLPEmbedder(params.vec_in_dim, self.hidden_size)
+        self.guidance_in = (
+            MLPEmbedder(in_dim=256, hidden_dim=self.hidden_size) if params.guidance_embed else nn.Identity()
+        )
+        self.txt_in = nn.Linear(params.context_in_dim, self.hidden_size)
+
+        self.double_blocks = nn.ModuleList(
+            [
+                DoubleStreamBlock(
+                    self.hidden_size,
+                    self.num_heads,
+                    mlp_ratio=params.mlp_ratio,
+                    qkv_bias=params.qkv_bias,
+                )
+                for _ in range(params.depth)
+            ]
+        )
+
+        self.single_blocks = nn.ModuleList(
+            [
+                SingleStreamBlock(self.hidden_size, self.num_heads, mlp_ratio=params.mlp_ratio)
+                for _ in range(params.depth_single_blocks)
+            ]
+        )
+
+        self.final_layer = LastLayer(self.hidden_size, 1, self.out_channels)
+
+    def forward(
+        self,
+        img: Tensor,
+        img_ids: Tensor,
+        txt: Tensor,
+        txt_ids: Tensor,
+        timesteps: Tensor,
+        y: Tensor,
+        guidance: Tensor | None = None,
+    ) -> Tensor:
+        if img.ndim != 3 or txt.ndim != 3:
+            raise ValueError("Input img and txt tensors must have 3 dimensions.")
+
+        # running on sequences img
+        img = self.img_in(img)
+        vec = self.time_in(timestep_embedding(timesteps, 256))
+        if self.params.guidance_embed:
+            if guidance is None:
+                raise ValueError("Didn't get guidance strength for guidance distilled model.")
+            vec = vec + self.guidance_in(timestep_embedding(guidance, 256))
+        vec = vec + self.vector_in(y)
+        txt = self.txt_in(txt)
+
+        ids = torch.cat((txt_ids, img_ids), dim=1)
+        pe = self.pe_embedder(ids)
+
+        for block in self.double_blocks:
+            img, txt = block(img=img, txt=txt, vec=vec, pe=pe)
+
+        img = torch.cat((txt, img), 1)
+        for block in self.single_blocks:
+            img = block(img, vec=vec, pe=pe)
+        img = img[:, txt.shape[1] :, ...]
+
+        img = self.final_layer(img, vec)  # (N, T, patch_size ** 2 * out_channels)
+        return img

invokeai/backend/flux/modules/autoencoder.py

@@ -0,0 +1,310 @@
+# Initially pulled from https://github.com/black-forest-labs/flux
+
+from dataclasses import dataclass
+
+import torch
+from einops import rearrange
+from torch import Tensor, nn
+
+
+@dataclass
+class AutoEncoderParams:
+    resolution: int
+    in_channels: int
+    ch: int
+    out_ch: int
+    ch_mult: list[int]
+    num_res_blocks: int
+    z_channels: int
+    scale_factor: float
+    shift_factor: float
+
+
+class AttnBlock(nn.Module):
+    def __init__(self, in_channels: int):
+        super().__init__()
+        self.in_channels = in_channels
+
+        self.norm = nn.GroupNorm(num_groups=32, num_channels=in_channels, eps=1e-6, affine=True)
+
+        self.q = nn.Conv2d(in_channels, in_channels, kernel_size=1)
+        self.k = nn.Conv2d(in_channels, in_channels, kernel_size=1)
+        self.v = nn.Conv2d(in_channels, in_channels, kernel_size=1)
+        self.proj_out = nn.Conv2d(in_channels, in_channels, kernel_size=1)
+
+    def attention(self, h_: Tensor) -> Tensor:
+        h_ = self.norm(h_)
+        q = self.q(h_)
+        k = self.k(h_)
+        v = self.v(h_)
+
+        b, c, h, w = q.shape
+        q = rearrange(q, "b c h w -> b 1 (h w) c").contiguous()
+        k = rearrange(k, "b c h w -> b 1 (h w) c").contiguous()
+        v = rearrange(v, "b c h w -> b 1 (h w) c").contiguous()
+        h_ = nn.functional.scaled_dot_product_attention(q, k, v)
+
+        return rearrange(h_, "b 1 (h w) c -> b c h w", h=h, w=w, c=c, b=b)
+
+    def forward(self, x: Tensor) -> Tensor:
+        return x + self.proj_out(self.attention(x))
+
+
+class ResnetBlock(nn.Module):
+    def __init__(self, in_channels: int, out_channels: int):
+        super().__init__()
+        self.in_channels = in_channels
+        out_channels = in_channels if out_channels is None else out_channels
+        self.out_channels = out_channels
+
+        self.norm1 = nn.GroupNorm(num_groups=32, num_channels=in_channels, eps=1e-6, affine=True)
+        self.conv1 = nn.Conv2d(in_channels, out_channels, kernel_size=3, stride=1, padding=1)
+        self.norm2 = nn.GroupNorm(num_groups=32, num_channels=out_channels, eps=1e-6, affine=True)
+        self.conv2 = nn.Conv2d(out_channels, out_channels, kernel_size=3, stride=1, padding=1)
+        if self.in_channels != self.out_channels:
+            self.nin_shortcut = nn.Conv2d(in_channels, out_channels, kernel_size=1, stride=1, padding=0)
+
+    def forward(self, x):
+        h = x
+        h = self.norm1(h)
+        h = torch.nn.functional.silu(h)
+        h = self.conv1(h)
+
+        h = self.norm2(h)
+        h = torch.nn.functional.silu(h)
+        h = self.conv2(h)
+
+        if self.in_channels != self.out_channels:
+            x = self.nin_shortcut(x)
+
+        return x + h
+
+
+class Downsample(nn.Module):
+    def __init__(self, in_channels: int):
+        super().__init__()
+        # no asymmetric padding in torch conv, must do it ourselves
+        self.conv = nn.Conv2d(in_channels, in_channels, kernel_size=3, stride=2, padding=0)
+
+    def forward(self, x: Tensor):
+        pad = (0, 1, 0, 1)
+        x = nn.functional.pad(x, pad, mode="constant", value=0)
+        x = self.conv(x)
+        return x
+
+
+class Upsample(nn.Module):
+    def __init__(self, in_channels: int):
+        super().__init__()
+        self.conv = nn.Conv2d(in_channels, in_channels, kernel_size=3, stride=1, padding=1)
+
+    def forward(self, x: Tensor):
+        x = nn.functional.interpolate(x, scale_factor=2.0, mode="nearest")
+        x = self.conv(x)
+        return x
+
+
+class Encoder(nn.Module):
+    def __init__(
+        self,
+        resolution: int,
+        in_channels: int,
+        ch: int,
+        ch_mult: list[int],
+        num_res_blocks: int,
+        z_channels: int,
+    ):
+        super().__init__()
+        self.ch = ch
+        self.num_resolutions = len(ch_mult)
+        self.num_res_blocks = num_res_blocks
+        self.resolution = resolution
+        self.in_channels = in_channels
+        # downsampling
+        self.conv_in = nn.Conv2d(in_channels, self.ch, kernel_size=3, stride=1, padding=1)
+
+        curr_res = resolution
+        in_ch_mult = (1,) + tuple(ch_mult)
+        self.in_ch_mult = in_ch_mult
+        self.down = nn.ModuleList()
+        block_in = self.ch
+        for i_level in range(self.num_resolutions):
+            block = nn.ModuleList()
+            attn = nn.ModuleList()
+            block_in = ch * in_ch_mult[i_level]
+            block_out = ch * ch_mult[i_level]
+            for _ in range(self.num_res_blocks):
+                block.append(ResnetBlock(in_channels=block_in, out_channels=block_out))
+                block_in = block_out
+            down = nn.Module()
+            down.block = block
+            down.attn = attn
+            if i_level != self.num_resolutions - 1:
+                down.downsample = Downsample(block_in)
+                curr_res = curr_res // 2
+            self.down.append(down)
+
+        # middle
+        self.mid = nn.Module()
+        self.mid.block_1 = ResnetBlock(in_channels=block_in, out_channels=block_in)
+        self.mid.attn_1 = AttnBlock(block_in)
+        self.mid.block_2 = ResnetBlock(in_channels=block_in, out_channels=block_in)
+
+        # end
+        self.norm_out = nn.GroupNorm(num_groups=32, num_channels=block_in, eps=1e-6, affine=True)
+        self.conv_out = nn.Conv2d(block_in, 2 * z_channels, kernel_size=3, stride=1, padding=1)
+
+    def forward(self, x: Tensor) -> Tensor:
+        # downsampling
+        hs = [self.conv_in(x)]
+        for i_level in range(self.num_resolutions):
+            for i_block in range(self.num_res_blocks):
+                h = self.down[i_level].block[i_block](hs[-1])
+                if len(self.down[i_level].attn) > 0:
+                    h = self.down[i_level].attn[i_block](h)
+                hs.append(h)
+            if i_level != self.num_resolutions - 1:
+                hs.append(self.down[i_level].downsample(hs[-1]))
+
+        # middle
+        h = hs[-1]
+        h = self.mid.block_1(h)
+        h = self.mid.attn_1(h)
+        h = self.mid.block_2(h)
+        # end
+        h = self.norm_out(h)
+        h = torch.nn.functional.silu(h)
+        h = self.conv_out(h)
+        return h
+
+
+class Decoder(nn.Module):
+    def __init__(
+        self,
+        ch: int,
+        out_ch: int,
+        ch_mult: list[int],
+        num_res_blocks: int,
+        in_channels: int,
+        resolution: int,
+        z_channels: int,
+    ):
+        super().__init__()
+        self.ch = ch
+        self.num_resolutions = len(ch_mult)
+        self.num_res_blocks = num_res_blocks
+        self.resolution = resolution
+        self.in_channels = in_channels
+        self.ffactor = 2 ** (self.num_resolutions - 1)
+
+        # compute in_ch_mult, block_in and curr_res at lowest res
+        block_in = ch * ch_mult[self.num_resolutions - 1]
+        curr_res = resolution // 2 ** (self.num_resolutions - 1)
+        self.z_shape = (1, z_channels, curr_res, curr_res)
+
+        # z to block_in
+        self.conv_in = nn.Conv2d(z_channels, block_in, kernel_size=3, stride=1, padding=1)
+
+        # middle
+        self.mid = nn.Module()
+        self.mid.block_1 = ResnetBlock(in_channels=block_in, out_channels=block_in)
+        self.mid.attn_1 = AttnBlock(block_in)
+        self.mid.block_2 = ResnetBlock(in_channels=block_in, out_channels=block_in)
+
+        # upsampling
+        self.up = nn.ModuleList()
+        for i_level in reversed(range(self.num_resolutions)):
+            block = nn.ModuleList()
+            attn = nn.ModuleList()
+            block_out = ch * ch_mult[i_level]
+            for _ in range(self.num_res_blocks + 1):
+                block.append(ResnetBlock(in_channels=block_in, out_channels=block_out))
+                block_in = block_out
+            up = nn.Module()
+            up.block = block
+            up.attn = attn
+            if i_level != 0:
+                up.upsample = Upsample(block_in)
+                curr_res = curr_res * 2
+            self.up.insert(0, up)  # prepend to get consistent order
+
+        # end
+        self.norm_out = nn.GroupNorm(num_groups=32, num_channels=block_in, eps=1e-6, affine=True)
+        self.conv_out = nn.Conv2d(block_in, out_ch, kernel_size=3, stride=1, padding=1)
+
+    def forward(self, z: Tensor) -> Tensor:
+        # z to block_in
+        h = self.conv_in(z)
+
+        # middle
+        h = self.mid.block_1(h)
+        h = self.mid.attn_1(h)
+        h = self.mid.block_2(h)
+
+        # upsampling
+        for i_level in reversed(range(self.num_resolutions)):
+            for i_block in range(self.num_res_blocks + 1):
+                h = self.up[i_level].block[i_block](h)
+                if len(self.up[i_level].attn) > 0:
+                    h = self.up[i_level].attn[i_block](h)
+            if i_level != 0:
+                h = self.up[i_level].upsample(h)
+
+        # end
+        h = self.norm_out(h)
+        h = torch.nn.functional.silu(h)
+        h = self.conv_out(h)
+        return h
+
+
+class DiagonalGaussian(nn.Module):
+    def __init__(self, sample: bool = True, chunk_dim: int = 1):
+        super().__init__()
+        self.sample = sample
+        self.chunk_dim = chunk_dim
+
+    def forward(self, z: Tensor) -> Tensor:
+        mean, logvar = torch.chunk(z, 2, dim=self.chunk_dim)
+        if self.sample:
+            std = torch.exp(0.5 * logvar)
+            return mean + std * torch.randn_like(mean)
+        else:
+            return mean
+
+
+class AutoEncoder(nn.Module):
+    def __init__(self, params: AutoEncoderParams):
+        super().__init__()
+        self.encoder = Encoder(
+            resolution=params.resolution,
+            in_channels=params.in_channels,
+            ch=params.ch,
+            ch_mult=params.ch_mult,
+            num_res_blocks=params.num_res_blocks,
+            z_channels=params.z_channels,
+        )
+        self.decoder = Decoder(
+            resolution=params.resolution,
+            in_channels=params.in_channels,
+            ch=params.ch,
+            out_ch=params.out_ch,
+            ch_mult=params.ch_mult,
+            num_res_blocks=params.num_res_blocks,
+            z_channels=params.z_channels,
+        )
+        self.reg = DiagonalGaussian()
+
+        self.scale_factor = params.scale_factor
+        self.shift_factor = params.shift_factor
+
+    def encode(self, x: Tensor) -> Tensor:
+        z = self.reg(self.encoder(x))
+        z = self.scale_factor * (z - self.shift_factor)
+        return z
+
+    def decode(self, z: Tensor) -> Tensor:
+        z = z / self.scale_factor + self.shift_factor
+        return self.decoder(z)
+
+    def forward(self, x: Tensor) -> Tensor:
+        return self.decode(self.encode(x))

invokeai/backend/flux/modules/conditioner.py

@@ -0,0 +1,33 @@
+# Initially pulled from https://github.com/black-forest-labs/flux
+
+from torch import Tensor, nn
+from transformers import PreTrainedModel, PreTrainedTokenizer
+
+
+class HFEncoder(nn.Module):
+    def __init__(self, encoder: PreTrainedModel, tokenizer: PreTrainedTokenizer, is_clip: bool, max_length: int):
+        super().__init__()
+        self.max_length = max_length
+        self.is_clip = is_clip
+        self.output_key = "pooler_output" if self.is_clip else "last_hidden_state"
+        self.tokenizer = tokenizer
+        self.hf_module = encoder
+        self.hf_module = self.hf_module.eval().requires_grad_(False)
+
+    def forward(self, text: list[str]) -> Tensor:
+        batch_encoding = self.tokenizer(
+            text,
+            truncation=True,
+            max_length=self.max_length,
+            return_length=False,
+            return_overflowing_tokens=False,
+            padding="max_length",
+            return_tensors="pt",
+        )
+
+        outputs = self.hf_module(
+            input_ids=batch_encoding["input_ids"].to(self.hf_module.device),
+            attention_mask=None,
+            output_hidden_states=False,
+        )
+        return outputs[self.output_key]

invokeai/backend/flux/modules/layers.py

@@ -0,0 +1,253 @@
+# Initially pulled from https://github.com/black-forest-labs/flux
+
+import math
+from dataclasses import dataclass
+
+import torch
+from einops import rearrange
+from torch import Tensor, nn
+
+from invokeai.backend.flux.math import attention, rope
+
+
+class EmbedND(nn.Module):
+    def __init__(self, dim: int, theta: int, axes_dim: list[int]):
+        super().__init__()
+        self.dim = dim
+        self.theta = theta
+        self.axes_dim = axes_dim
+
+    def forward(self, ids: Tensor) -> Tensor:
+        n_axes = ids.shape[-1]
+        emb = torch.cat(
+            [rope(ids[..., i], self.axes_dim[i], self.theta) for i in range(n_axes)],
+            dim=-3,
+        )
+
+        return emb.unsqueeze(1)
+
+
+def timestep_embedding(t: Tensor, dim, max_period=10000, time_factor: float = 1000.0):
+    """
+    Create sinusoidal timestep embeddings.
+    :param t: a 1-D Tensor of N indices, one per batch element.
+                      These may be fractional.
+    :param dim: the dimension of the output.
+    :param max_period: controls the minimum frequency of the embeddings.
+    :return: an (N, D) Tensor of positional embeddings.
+    """
+    t = time_factor * t
+    half = dim // 2
+    freqs = torch.exp(-math.log(max_period) * torch.arange(start=0, end=half, dtype=torch.float32) / half).to(t.device)
+
+    args = t[:, None].float() * freqs[None]
+    embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1)
+    if dim % 2:
+        embedding = torch.cat([embedding, torch.zeros_like(embedding[:, :1])], dim=-1)
+    if torch.is_floating_point(t):
+        embedding = embedding.to(t)
+    return embedding
+
+
+class MLPEmbedder(nn.Module):
+    def __init__(self, in_dim: int, hidden_dim: int):
+        super().__init__()
+        self.in_layer = nn.Linear(in_dim, hidden_dim, bias=True)
+        self.silu = nn.SiLU()
+        self.out_layer = nn.Linear(hidden_dim, hidden_dim, bias=True)
+
+    def forward(self, x: Tensor) -> Tensor:
+        return self.out_layer(self.silu(self.in_layer(x)))
+
+
+class RMSNorm(torch.nn.Module):
+    def __init__(self, dim: int):
+        super().__init__()
+        self.scale = nn.Parameter(torch.ones(dim))
+
+    def forward(self, x: Tensor):
+        x_dtype = x.dtype
+        x = x.float()
+        rrms = torch.rsqrt(torch.mean(x**2, dim=-1, keepdim=True) + 1e-6)
+        return (x * rrms).to(dtype=x_dtype) * self.scale
+
+
+class QKNorm(torch.nn.Module):
+    def __init__(self, dim: int):
+        super().__init__()
+        self.query_norm = RMSNorm(dim)
+        self.key_norm = RMSNorm(dim)
+
+    def forward(self, q: Tensor, k: Tensor, v: Tensor) -> tuple[Tensor, Tensor]:
+        q = self.query_norm(q)
+        k = self.key_norm(k)
+        return q.to(v), k.to(v)
+
+
+class SelfAttention(nn.Module):
+    def __init__(self, dim: int, num_heads: int = 8, qkv_bias: bool = False):
+        super().__init__()
+        self.num_heads = num_heads
+        head_dim = dim // num_heads
+
+        self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)
+        self.norm = QKNorm(head_dim)
+        self.proj = nn.Linear(dim, dim)
+
+    def forward(self, x: Tensor, pe: Tensor) -> Tensor:
+        qkv = self.qkv(x)
+        q, k, v = rearrange(qkv, "B L (K H D) -> K B H L D", K=3, H=self.num_heads)
+        q, k = self.norm(q, k, v)
+        x = attention(q, k, v, pe=pe)
+        x = self.proj(x)
+        return x
+
+
+@dataclass
+class ModulationOut:
+    shift: Tensor
+    scale: Tensor
+    gate: Tensor
+
+
+class Modulation(nn.Module):
+    def __init__(self, dim: int, double: bool):
+        super().__init__()
+        self.is_double = double
+        self.multiplier = 6 if double else 3
+        self.lin = nn.Linear(dim, self.multiplier * dim, bias=True)
+
+    def forward(self, vec: Tensor) -> tuple[ModulationOut, ModulationOut | None]:
+        out = self.lin(nn.functional.silu(vec))[:, None, :].chunk(self.multiplier, dim=-1)
+
+        return (
+            ModulationOut(*out[:3]),
+            ModulationOut(*out[3:]) if self.is_double else None,
+        )
+
+
+class DoubleStreamBlock(nn.Module):
+    def __init__(self, hidden_size: int, num_heads: int, mlp_ratio: float, qkv_bias: bool = False):
+        super().__init__()
+
+        mlp_hidden_dim = int(hidden_size * mlp_ratio)
+        self.num_heads = num_heads
+        self.hidden_size = hidden_size
+        self.img_mod = Modulation(hidden_size, double=True)
+        self.img_norm1 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
+        self.img_attn = SelfAttention(dim=hidden_size, num_heads=num_heads, qkv_bias=qkv_bias)
+
+        self.img_norm2 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
+        self.img_mlp = nn.Sequential(
+            nn.Linear(hidden_size, mlp_hidden_dim, bias=True),
+            nn.GELU(approximate="tanh"),
+            nn.Linear(mlp_hidden_dim, hidden_size, bias=True),
+        )
+
+        self.txt_mod = Modulation(hidden_size, double=True)
+        self.txt_norm1 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
+        self.txt_attn = SelfAttention(dim=hidden_size, num_heads=num_heads, qkv_bias=qkv_bias)
+
+        self.txt_norm2 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
+        self.txt_mlp = nn.Sequential(
+            nn.Linear(hidden_size, mlp_hidden_dim, bias=True),
+            nn.GELU(approximate="tanh"),
+            nn.Linear(mlp_hidden_dim, hidden_size, bias=True),
+        )
+
+    def forward(self, img: Tensor, txt: Tensor, vec: Tensor, pe: Tensor) -> tuple[Tensor, Tensor]:
+        img_mod1, img_mod2 = self.img_mod(vec)
+        txt_mod1, txt_mod2 = self.txt_mod(vec)
+
+        # prepare image for attention
+        img_modulated = self.img_norm1(img)
+        img_modulated = (1 + img_mod1.scale) * img_modulated + img_mod1.shift
+        img_qkv = self.img_attn.qkv(img_modulated)
+        img_q, img_k, img_v = rearrange(img_qkv, "B L (K H D) -> K B H L D", K=3, H=self.num_heads)
+        img_q, img_k = self.img_attn.norm(img_q, img_k, img_v)
+
+        # prepare txt for attention
+        txt_modulated = self.txt_norm1(txt)
+        txt_modulated = (1 + txt_mod1.scale) * txt_modulated + txt_mod1.shift
+        txt_qkv = self.txt_attn.qkv(txt_modulated)
+        txt_q, txt_k, txt_v = rearrange(txt_qkv, "B L (K H D) -> K B H L D", K=3, H=self.num_heads)
+        txt_q, txt_k = self.txt_attn.norm(txt_q, txt_k, txt_v)
+
+        # run actual attention
+        q = torch.cat((txt_q, img_q), dim=2)
+        k = torch.cat((txt_k, img_k), dim=2)
+        v = torch.cat((txt_v, img_v), dim=2)
+
+        attn = attention(q, k, v, pe=pe)
+        txt_attn, img_attn = attn[:, : txt.shape[1]], attn[:, txt.shape[1] :]
+
+        # calculate the img bloks
+        img = img + img_mod1.gate * self.img_attn.proj(img_attn)
+        img = img + img_mod2.gate * self.img_mlp((1 + img_mod2.scale) * self.img_norm2(img) + img_mod2.shift)
+
+        # calculate the txt bloks
+        txt = txt + txt_mod1.gate * self.txt_attn.proj(txt_attn)
+        txt = txt + txt_mod2.gate * self.txt_mlp((1 + txt_mod2.scale) * self.txt_norm2(txt) + txt_mod2.shift)
+        return img, txt
+
+
+class SingleStreamBlock(nn.Module):
+    """
+    A DiT block with parallel linear layers as described in
+    https://arxiv.org/abs/2302.05442 and adapted modulation interface.
+    """
+
+    def __init__(
+        self,
+        hidden_size: int,
+        num_heads: int,
+        mlp_ratio: float = 4.0,
+        qk_scale: float | None = None,
+    ):
+        super().__init__()
+        self.hidden_dim = hidden_size
+        self.num_heads = num_heads
+        head_dim = hidden_size // num_heads
+        self.scale = qk_scale or head_dim**-0.5
+
+        self.mlp_hidden_dim = int(hidden_size * mlp_ratio)
+        # qkv and mlp_in
+        self.linear1 = nn.Linear(hidden_size, hidden_size * 3 + self.mlp_hidden_dim)
+        # proj and mlp_out
+        self.linear2 = nn.Linear(hidden_size + self.mlp_hidden_dim, hidden_size)
+
+        self.norm = QKNorm(head_dim)
+
+        self.hidden_size = hidden_size
+        self.pre_norm = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
+
+        self.mlp_act = nn.GELU(approximate="tanh")
+        self.modulation = Modulation(hidden_size, double=False)
+
+    def forward(self, x: Tensor, vec: Tensor, pe: Tensor) -> Tensor:
+        mod, _ = self.modulation(vec)
+        x_mod = (1 + mod.scale) * self.pre_norm(x) + mod.shift
+        qkv, mlp = torch.split(self.linear1(x_mod), [3 * self.hidden_size, self.mlp_hidden_dim], dim=-1)
+
+        q, k, v = rearrange(qkv, "B L (K H D) -> K B H L D", K=3, H=self.num_heads)
+        q, k = self.norm(q, k, v)
+
+        # compute attention
+        attn = attention(q, k, v, pe=pe)
+        # compute activation in mlp stream, cat again and run second linear layer
+        output = self.linear2(torch.cat((attn, self.mlp_act(mlp)), 2))
+        return x + mod.gate * output
+
+
+class LastLayer(nn.Module):
+    def __init__(self, hidden_size: int, patch_size: int, out_channels: int):
+        super().__init__()
+        self.norm_final = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
+        self.linear = nn.Linear(hidden_size, patch_size * patch_size * out_channels, bias=True)
+        self.adaLN_modulation = nn.Sequential(nn.SiLU(), nn.Linear(hidden_size, 2 * hidden_size, bias=True))
+
+    def forward(self, x: Tensor, vec: Tensor) -> Tensor:
+        shift, scale = self.adaLN_modulation(vec).chunk(2, dim=1)
+        x = (1 + scale[:, None, :]) * self.norm_final(x) + shift[:, None, :]
+        x = self.linear(x)
+        return x

invokeai/backend/flux/sampling.py

@@ -0,0 +1,167 @@
+# Initially pulled from https://github.com/black-forest-labs/flux
+
+import math
+from typing import Callable
+
+import torch
+from einops import rearrange, repeat
+from torch import Tensor
+from tqdm import tqdm
+
+from invokeai.backend.flux.model import Flux
+from invokeai.backend.flux.modules.conditioner import HFEncoder
+
+
+def get_noise(
+    num_samples: int,
+    height: int,
+    width: int,
+    device: torch.device,
+    dtype: torch.dtype,
+    seed: int,
+):
+    # We always generate noise on the same device and dtype then cast to ensure consistency across devices/dtypes.
+    rand_device = "cpu"
+    rand_dtype = torch.float16
+    return torch.randn(
+        num_samples,
+        16,
+        # allow for packing
+        2 * math.ceil(height / 16),
+        2 * math.ceil(width / 16),
+        device=rand_device,
+        dtype=rand_dtype,
+        generator=torch.Generator(device=rand_device).manual_seed(seed),
+    ).to(device=device, dtype=dtype)
+
+
+def prepare(t5: HFEncoder, clip: HFEncoder, img: Tensor, prompt: str | list[str]) -> dict[str, Tensor]:
+    bs, c, h, w = img.shape
+    if bs == 1 and not isinstance(prompt, str):
+        bs = len(prompt)
+
+    img = rearrange(img, "b c (h ph) (w pw) -> b (h w) (c ph pw)", ph=2, pw=2)
+    if img.shape[0] == 1 and bs > 1:
+        img = repeat(img, "1 ... -> bs ...", bs=bs)
+
+    img_ids = torch.zeros(h // 2, w // 2, 3)
+    img_ids[..., 1] = img_ids[..., 1] + torch.arange(h // 2)[:, None]
+    img_ids[..., 2] = img_ids[..., 2] + torch.arange(w // 2)[None, :]
+    img_ids = repeat(img_ids, "h w c -> b (h w) c", b=bs)
+
+    if isinstance(prompt, str):
+        prompt = [prompt]
+    txt = t5(prompt)
+    if txt.shape[0] == 1 and bs > 1:
+        txt = repeat(txt, "1 ... -> bs ...", bs=bs)
+    txt_ids = torch.zeros(bs, txt.shape[1], 3)
+
+    vec = clip(prompt)
+    if vec.shape[0] == 1 and bs > 1:
+        vec = repeat(vec, "1 ... -> bs ...", bs=bs)
+
+    return {
+        "img": img,
+        "img_ids": img_ids.to(img.device),
+        "txt": txt.to(img.device),
+        "txt_ids": txt_ids.to(img.device),
+        "vec": vec.to(img.device),
+    }
+
+
+def time_shift(mu: float, sigma: float, t: Tensor):
+    return math.exp(mu) / (math.exp(mu) + (1 / t - 1) ** sigma)
+
+
+def get_lin_function(x1: float = 256, y1: float = 0.5, x2: float = 4096, y2: float = 1.15) -> Callable[[float], float]:
+    m = (y2 - y1) / (x2 - x1)
+    b = y1 - m * x1
+    return lambda x: m * x + b
+
+
+def get_schedule(
+    num_steps: int,
+    image_seq_len: int,
+    base_shift: float = 0.5,
+    max_shift: float = 1.15,
+    shift: bool = True,
+) -> list[float]:
+    # extra step for zero
+    timesteps = torch.linspace(1, 0, num_steps + 1)
+
+    # shifting the schedule to favor high timesteps for higher signal images
+    if shift:
+        # eastimate mu based on linear estimation between two points
+        mu = get_lin_function(y1=base_shift, y2=max_shift)(image_seq_len)
+        timesteps = time_shift(mu, 1.0, timesteps)
+
+    return timesteps.tolist()
+
+
+def denoise(
+    model: Flux,
+    # model input
+    img: Tensor,
+    img_ids: Tensor,
+    txt: Tensor,
+    txt_ids: Tensor,
+    vec: Tensor,
+    # sampling parameters
+    timesteps: list[float],
+    step_callback: Callable[[], None],
+    guidance: float = 4.0,
+):
+    # guidance_vec is ignored for schnell.
+    guidance_vec = torch.full((img.shape[0],), guidance, device=img.device, dtype=img.dtype)
+    for t_curr, t_prev in tqdm(list(zip(timesteps[:-1], timesteps[1:], strict=True))):
+        t_vec = torch.full((img.shape[0],), t_curr, dtype=img.dtype, device=img.device)
+        pred = model(
+            img=img,
+            img_ids=img_ids,
+            txt=txt,
+            txt_ids=txt_ids,
+            y=vec,
+            timesteps=t_vec,
+            guidance=guidance_vec,
+        )
+
+        img = img + (t_prev - t_curr) * pred
+        step_callback()
+
+    return img
+
+
+def unpack(x: Tensor, height: int, width: int) -> Tensor:
+    return rearrange(
+        x,
+        "b (h w) (c ph pw) -> b c (h ph) (w pw)",
+        h=math.ceil(height / 16),
+        w=math.ceil(width / 16),
+        ph=2,
+        pw=2,
+    )
+
+
+def prepare_latent_img_patches(latent_img: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+    """Convert an input image in latent space to patches for diffusion.
+
+    This implementation was extracted from:
+    https://github.com/black-forest-labs/flux/blob/c00d7c60b085fce8058b9df845e036090873f2ce/src/flux/sampling.py#L32
+
+    Returns:
+        tuple[Tensor, Tensor]: (img, img_ids), as defined in the original flux repo.
+    """
+    bs, c, h, w = latent_img.shape
+
+    # Pixel unshuffle with a scale of 2, and flatten the height/width dimensions to get an array of patches.
+    img = rearrange(latent_img, "b c (h ph) (w pw) -> b (h w) (c ph pw)", ph=2, pw=2)
+    if img.shape[0] == 1 and bs > 1:
+        img = repeat(img, "1 ... -> bs ...", bs=bs)
+
+    # Generate patch position ids.
+    img_ids = torch.zeros(h // 2, w // 2, 3, device=img.device, dtype=img.dtype)
+    img_ids[..., 1] = img_ids[..., 1] + torch.arange(h // 2, device=img.device, dtype=img.dtype)[:, None]
+    img_ids[..., 2] = img_ids[..., 2] + torch.arange(w // 2, device=img.device, dtype=img.dtype)[None, :]
+    img_ids = repeat(img_ids, "h w c -> b (h w) c", b=bs)
+
+    return img, img_ids

invokeai/backend/flux/util.py

@@ -0,0 +1,71 @@
+# Initially pulled from https://github.com/black-forest-labs/flux
+
+from dataclasses import dataclass
+from typing import Dict, Literal
+
+from invokeai.backend.flux.model import FluxParams
+from invokeai.backend.flux.modules.autoencoder import AutoEncoderParams
+
+
+@dataclass
+class ModelSpec:
+    params: FluxParams
+    ae_params: AutoEncoderParams
+    ckpt_path: str | None
+    ae_path: str | None
+    repo_id: str | None
+    repo_flow: str | None
+    repo_ae: str | None
+
+
+max_seq_lengths: Dict[str, Literal[256, 512]] = {
+    "flux-dev": 512,
+    "flux-schnell": 256,
+}
+
+
+ae_params = {
+    "flux": AutoEncoderParams(
+        resolution=256,
+        in_channels=3,
+        ch=128,
+        out_ch=3,
+        ch_mult=[1, 2, 4, 4],
+        num_res_blocks=2,
+        z_channels=16,
+        scale_factor=0.3611,
+        shift_factor=0.1159,
+    )
+}
+
+
+params = {
+    "flux-dev": FluxParams(
+        in_channels=64,
+        vec_in_dim=768,
+        context_in_dim=4096,
+        hidden_size=3072,
+        mlp_ratio=4.0,
+        num_heads=24,
+        depth=19,
+        depth_single_blocks=38,
+        axes_dim=[16, 56, 56],
+        theta=10_000,
+        qkv_bias=True,
+        guidance_embed=True,
+    ),
+    "flux-schnell": FluxParams(
+        in_channels=64,
+        vec_in_dim=768,
+        context_in_dim=4096,
+        hidden_size=3072,
+        mlp_ratio=4.0,
+        num_heads=24,
+        depth=19,
+        depth_single_blocks=38,
+        axes_dim=[16, 56, 56],
+        theta=10_000,
+        qkv_bias=True,
+        guidance_embed=False,
+    ),
+}

invokeai/backend/model_manager/config.py

@@ -52,6 +52,7 @@ class BaseModelType(str, Enum):
     StableDiffusion2 = "sd-2"
     StableDiffusionXL = "sdxl"
     StableDiffusionXLRefiner = "sdxl-refiner"
+    Flux = "flux"
     # Kandinsky2_1 = "kandinsky-2.1"
 
 
@@ -66,7 +67,9 @@ class ModelType(str, Enum):
     TextualInversion = "embedding"
     IPAdapter = "ip_adapter"
     CLIPVision = "clip_vision"
+    CLIPEmbed = "clip_embed"
     T2IAdapter = "t2i_adapter"
+    T5Encoder = "t5_encoder"
     SpandrelImageToImage = "spandrel_image_to_image"
 
 
@@ -74,6 +77,7 @@ class SubModelType(str, Enum):
     """Submodel type."""
 
     UNet = "unet"
+    Transformer = "transformer"
     TextEncoder = "text_encoder"
     TextEncoder2 = "text_encoder_2"
     Tokenizer = "tokenizer"
@@ -104,6 +108,9 @@ class ModelFormat(str, Enum):
     EmbeddingFile = "embedding_file"
     EmbeddingFolder = "embedding_folder"
     InvokeAI = "invokeai"
+    T5Encoder = "t5_encoder"
+    BnbQuantizedLlmInt8b = "bnb_quantized_int8b"
+    BnbQuantizednf4b = "bnb_quantized_nf4b"
 
 
 class SchedulerPredictionType(str, Enum):
@@ -186,7 +193,9 @@ class ModelConfigBase(BaseModel):
 class CheckpointConfigBase(ModelConfigBase):
     """Model config for checkpoint-style models."""
 
-    format: Literal[ModelFormat.Checkpoint] = ModelFormat.Checkpoint
+    format: Literal[ModelFormat.Checkpoint, ModelFormat.BnbQuantizednf4b] = Field(
+        description="Format of the provided checkpoint model", default=ModelFormat.Checkpoint
+    )
     config_path: str = Field(description="path to the checkpoint model config file")
     converted_at: Optional[float] = Field(
         description="When this model was last converted to diffusers", default_factory=time.time
@@ -205,6 +214,26 @@ class LoRAConfigBase(ModelConfigBase):
     trigger_phrases: Optional[set[str]] = Field(description="Set of trigger phrases for this model", default=None)
 
 
+class T5EncoderConfigBase(ModelConfigBase):
+    type: Literal[ModelType.T5Encoder] = ModelType.T5Encoder
+
+
+class T5EncoderConfig(T5EncoderConfigBase):
+    format: Literal[ModelFormat.T5Encoder] = ModelFormat.T5Encoder
+
+    @staticmethod
+    def get_tag() -> Tag:
+        return Tag(f"{ModelType.T5Encoder.value}.{ModelFormat.T5Encoder.value}")
+
+
+class T5EncoderBnbQuantizedLlmInt8bConfig(T5EncoderConfigBase):
+    format: Literal[ModelFormat.BnbQuantizedLlmInt8b] = ModelFormat.BnbQuantizedLlmInt8b
+
+    @staticmethod
+    def get_tag() -> Tag:
+        return Tag(f"{ModelType.T5Encoder.value}.{ModelFormat.BnbQuantizedLlmInt8b.value}")
+
+
 class LoRALyCORISConfig(LoRAConfigBase):
     """Model config for LoRA/Lycoris models."""
 
@@ -229,7 +258,6 @@ class VAECheckpointConfig(CheckpointConfigBase):
     """Model config for standalone VAE models."""
 
     type: Literal[ModelType.VAE] = ModelType.VAE
-    format: Literal[ModelFormat.Checkpoint] = ModelFormat.Checkpoint
 
     @staticmethod
     def get_tag() -> Tag:
@@ -268,7 +296,6 @@ class ControlNetCheckpointConfig(CheckpointConfigBase, ControlAdapterConfigBase)
     """Model config for ControlNet models (diffusers version)."""
 
     type: Literal[ModelType.ControlNet] = ModelType.ControlNet
-    format: Literal[ModelFormat.Checkpoint] = ModelFormat.Checkpoint
 
     @staticmethod
     def get_tag() -> Tag:
@@ -317,6 +344,21 @@ class MainCheckpointConfig(CheckpointConfigBase, MainConfigBase):
         return Tag(f"{ModelType.Main.value}.{ModelFormat.Checkpoint.value}")
 
 
+class MainBnbQuantized4bCheckpointConfig(CheckpointConfigBase, MainConfigBase):
+    """Model config for main checkpoint models."""
+
+    prediction_type: SchedulerPredictionType = SchedulerPredictionType.Epsilon
+    upcast_attention: bool = False
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.format = ModelFormat.BnbQuantizednf4b
+
+    @staticmethod
+    def get_tag() -> Tag:
+        return Tag(f"{ModelType.Main.value}.{ModelFormat.BnbQuantizednf4b.value}")
+
+
 class MainDiffusersConfig(DiffusersConfigBase, MainConfigBase):
     """Model config for main diffusers models."""
 
@@ -350,6 +392,17 @@ class IPAdapterCheckpointConfig(IPAdapterBaseConfig):
         return Tag(f"{ModelType.IPAdapter.value}.{ModelFormat.Checkpoint.value}")
 
 
+class CLIPEmbedDiffusersConfig(DiffusersConfigBase):
+    """Model config for Clip Embeddings."""
+
+    type: Literal[ModelType.CLIPEmbed] = ModelType.CLIPEmbed
+    format: Literal[ModelFormat.Diffusers] = ModelFormat.Diffusers
+
+    @staticmethod
+    def get_tag() -> Tag:
+        return Tag(f"{ModelType.CLIPEmbed.value}.{ModelFormat.Diffusers.value}")
+
+
 class CLIPVisionDiffusersConfig(DiffusersConfigBase):
     """Model config for CLIPVision."""
 
@@ -408,12 +461,15 @@ AnyModelConfig = Annotated[
     Union[
         Annotated[MainDiffusersConfig, MainDiffusersConfig.get_tag()],
         Annotated[MainCheckpointConfig, MainCheckpointConfig.get_tag()],
+        Annotated[MainBnbQuantized4bCheckpointConfig, MainBnbQuantized4bCheckpointConfig.get_tag()],
         Annotated[VAEDiffusersConfig, VAEDiffusersConfig.get_tag()],
         Annotated[VAECheckpointConfig, VAECheckpointConfig.get_tag()],
         Annotated[ControlNetDiffusersConfig, ControlNetDiffusersConfig.get_tag()],
         Annotated[ControlNetCheckpointConfig, ControlNetCheckpointConfig.get_tag()],
         Annotated[LoRALyCORISConfig, LoRALyCORISConfig.get_tag()],
         Annotated[LoRADiffusersConfig, LoRADiffusersConfig.get_tag()],
+        Annotated[T5EncoderConfig, T5EncoderConfig.get_tag()],
+        Annotated[T5EncoderBnbQuantizedLlmInt8bConfig, T5EncoderBnbQuantizedLlmInt8bConfig.get_tag()],
         Annotated[TextualInversionFileConfig, TextualInversionFileConfig.get_tag()],
         Annotated[TextualInversionFolderConfig, TextualInversionFolderConfig.get_tag()],
         Annotated[IPAdapterInvokeAIConfig, IPAdapterInvokeAIConfig.get_tag()],
@@ -421,6 +477,7 @@ AnyModelConfig = Annotated[
         Annotated[T2IAdapterConfig, T2IAdapterConfig.get_tag()],
         Annotated[SpandrelImageToImageConfig, SpandrelImageToImageConfig.get_tag()],
         Annotated[CLIPVisionDiffusersConfig, CLIPVisionDiffusersConfig.get_tag()],
+        Annotated[CLIPEmbedDiffusersConfig, CLIPEmbedDiffusersConfig.get_tag()],
     ],
     Discriminator(get_model_discriminator_value),
 ]

invokeai/backend/model_manager/load/load_default.py

@@ -72,6 +72,7 @@ class ModelLoader(ModelLoaderBase):
             pass
 
         config.path = str(self._get_model_path(config))
+        self._ram_cache.make_room(self.get_size_fs(config, Path(config.path), submodel_type))
         loaded_model = self._load_model(config, submodel_type)
 
         self._ram_cache.put(

invokeai/backend/model_manager/load/model_cache/model_cache_base.py

@@ -193,15 +193,6 @@ class ModelCacheBase(ABC, Generic[T]):
         """
         pass
 
-    @abstractmethod
-    def exists(
-        self,
-        key: str,
-        submodel_type: Optional[SubModelType] = None,
-    ) -> bool:
-        """Return true if the model identified by key and submodel_type is in the cache."""
-        pass
-
     @abstractmethod
     def cache_size(self) -> int:
         """Get the total size of the models currently cached."""

invokeai/backend/model_manager/load/model_cache/model_cache_default.py

@@ -1,22 +1,6 @@
 # Copyright (c) 2024 Lincoln D. Stein and the InvokeAI Development team
 # TODO: Add Stalker's proper name to copyright
-"""
-Manage a RAM cache of diffusion/transformer models for fast switching.
-They are moved between GPU VRAM and CPU RAM as necessary. If the cache
-grows larger than a preset maximum, then the least recently used
-model will be cleared and (re)loaded from disk when next needed.
-
-The cache returns context manager generators designed to load the
-model into the GPU within the context, and unload outside the
-context. Use like this:
-
-   cache = ModelCache(max_cache_size=7.5)
-   with cache.get_model('runwayml/stable-diffusion-1-5') as SD1,
-          cache.get_model('stabilityai/stable-diffusion-2') as SD2:
-       do_something_in_GPU(SD1,SD2)
-
-
-"""
+""" """
 
 import gc
 import math
@@ -40,45 +24,64 @@ from invokeai.backend.model_manager.load.model_util import calc_model_size_by_da
 from invokeai.backend.util.devices import TorchDevice
 from invokeai.backend.util.logging import InvokeAILogger
 
-# Maximum size of the cache, in gigs
-# Default is roughly enough to hold three fp16 diffusers models in RAM simultaneously
-DEFAULT_MAX_CACHE_SIZE = 6.0
-
-# amount of GPU memory to hold in reserve for use by generations (GB)
-DEFAULT_MAX_VRAM_CACHE_SIZE = 2.75
-
-# actual size of a gig
-GIG = 1073741824
+# Size of a GB in bytes.
+GB = 2**30
 
 # Size of a MB in bytes.
 MB = 2**20
 
 
 class ModelCache(ModelCacheBase[AnyModel]):
-    """Implementation of ModelCacheBase."""
+    """A cache for managing models in memory.
+
+    The cache is based on two levels of model storage:
+    - execution_device: The device where most models are executed (typically "cuda", "mps", or "cpu").
+    - storage_device: The device where models are offloaded when not in active use (typically "cpu").
+
+    The model cache is based on the following assumptions:
+    - storage_device_mem_size > execution_device_mem_size
+    - disk_to_storage_device_transfer_time >> storage_device_to_execution_device_transfer_time
+
+    A copy of all models in the cache is always kept on the storage_device. A subset of the models also have a copy on
+    the execution_device.
+
+    Models are moved between the storage_device and the execution_device as necessary. Cache size limits are enforced
+    on both the storage_device and the execution_device. The execution_device cache uses a smallest-first offload
+    policy. The storage_device cache uses a least-recently-used (LRU) offload policy.
+
+    Note: Neither of these offload policies has really been compared against alternatives. It's likely that different
+    policies would be better, although the optimal policies are likely heavily dependent on usage patterns and HW
+    configuration.
+
+    The cache returns context manager generators designed to load the model into the execution device (often GPU) within
+    the context, and unload outside the context.
+
+    Example usage:
+    ```
+    cache = ModelCache(max_cache_size=7.5, max_vram_cache_size=6.0)
+    with cache.get_model('runwayml/stable-diffusion-1-5') as SD1:
+        do_something_on_gpu(SD1)
+    ```
+    """
 
     def __init__(
         self,
-        max_cache_size: float = DEFAULT_MAX_CACHE_SIZE,
-        max_vram_cache_size: float = DEFAULT_MAX_VRAM_CACHE_SIZE,
+        max_cache_size: float,
+        max_vram_cache_size: float,
         execution_device: torch.device = torch.device("cuda"),
         storage_device: torch.device = torch.device("cpu"),
-        precision: torch.dtype = torch.float16,
-        sequential_offload: bool = False,
         lazy_offloading: bool = True,
-        sha_chunksize: int = 16777216,
         log_memory_usage: bool = False,
         logger: Optional[Logger] = None,
     ):
         """
         Initialize the model RAM cache.
 
-        :param max_cache_size: Maximum size of the RAM cache [6.0 GB]
+        :param max_cache_size: Maximum size of the storage_device cache in GBs.
+        :param max_vram_cache_size: Maximum size of the execution_device cache in GBs.
         :param execution_device: Torch device to load active model into [torch.device('cuda')]
         :param storage_device: Torch device to save inactive model in [torch.device('cpu')]
-        :param precision: Precision for loaded models [torch.float16]
-        :param lazy_offloading: Keep model in VRAM until another model needs to be loaded
-        :param sequential_offload: Conserve VRAM by loading and unloading each stage of the pipeline sequentially
+        :param lazy_offloading: Keep model in VRAM until another model needs to be loaded.
         :param log_memory_usage: If True, a memory snapshot will be captured before and after every model cache
             operation, and the result will be logged (at debug level). There is a time cost to capturing the memory
             snapshots, so it is recommended to disable this feature unless you are actively inspecting the model cache's
@@ -86,7 +89,6 @@ class ModelCache(ModelCacheBase[AnyModel]):
         """
         # allow lazy offloading only when vram cache enabled
         self._lazy_offloading = lazy_offloading and max_vram_cache_size > 0
-        self._precision: torch.dtype = precision
         self._max_cache_size: float = max_cache_size
         self._max_vram_cache_size: float = max_vram_cache_size
         self._execution_device: torch.device = execution_device
@@ -145,15 +147,6 @@ class ModelCache(ModelCacheBase[AnyModel]):
             total += cache_record.size
         return total
 
-    def exists(
-        self,
-        key: str,
-        submodel_type: Optional[SubModelType] = None,
-    ) -> bool:
-        """Return true if the model identified by key and submodel_type is in the cache."""
-        key = self._make_cache_key(key, submodel_type)
-        return key in self._cached_models
-
     def put(
         self,
         key: str,
@@ -203,7 +196,7 @@ class ModelCache(ModelCacheBase[AnyModel]):
         # more stats
         if self.stats:
             stats_name = stats_name or key
-            self.stats.cache_size = int(self._max_cache_size * GIG)
+            self.stats.cache_size = int(self._max_cache_size * GB)
             self.stats.high_watermark = max(self.stats.high_watermark, self.cache_size())
             self.stats.in_cache = len(self._cached_models)
             self.stats.loaded_model_sizes[stats_name] = max(
@@ -231,10 +224,13 @@ class ModelCache(ModelCacheBase[AnyModel]):
             return model_key
 
     def offload_unlocked_models(self, size_required: int) -> None:
-        """Move any unused models from VRAM."""
-        reserved = self._max_vram_cache_size * GIG
+        """Offload models from the execution_device to make room for size_required.
+
+        :param size_required: The amount of space to clear in the execution_device cache, in bytes.
+        """
+        reserved = self._max_vram_cache_size * GB
         vram_in_use = torch.cuda.memory_allocated() + size_required
-        self.logger.debug(f"{(vram_in_use/GIG):.2f}GB VRAM needed for models; max allowed={(reserved/GIG):.2f}GB")
+        self.logger.debug(f"{(vram_in_use/GB):.2f}GB VRAM needed for models; max allowed={(reserved/GB):.2f}GB")
         for _, cache_entry in sorted(self._cached_models.items(), key=lambda x: x[1].size):
             if vram_in_use <= reserved:
                 break
@@ -245,7 +241,7 @@ class ModelCache(ModelCacheBase[AnyModel]):
                 cache_entry.loaded = False
                 vram_in_use = torch.cuda.memory_allocated() + size_required
                 self.logger.debug(
-                    f"Removing {cache_entry.key} from VRAM to free {(cache_entry.size/GIG):.2f}GB; vram free = {(torch.cuda.memory_allocated()/GIG):.2f}GB"
+                    f"Removing {cache_entry.key} from VRAM to free {(cache_entry.size/GB):.2f}GB; vram free = {(torch.cuda.memory_allocated()/GB):.2f}GB"
                 )
 
         TorchDevice.empty_cache()
@@ -303,7 +299,7 @@ class ModelCache(ModelCacheBase[AnyModel]):
         self.logger.debug(
             f"Moved model '{cache_entry.key}' from {source_device} to"
             f" {target_device} in {(end_model_to_time-start_model_to_time):.2f}s."
-            f"Estimated model size: {(cache_entry.size/GIG):.3f} GB."
+            f"Estimated model size: {(cache_entry.size/GB):.3f} GB."
             f"{get_pretty_snapshot_diff(snapshot_before, snapshot_after)}"
         )
 
@@ -326,14 +322,14 @@ class ModelCache(ModelCacheBase[AnyModel]):
                     f"Moving model '{cache_entry.key}' from {source_device} to"
                     f" {target_device} caused an unexpected change in VRAM usage. The model's"
                     " estimated size may be incorrect. Estimated model size:"
-                    f" {(cache_entry.size/GIG):.3f} GB.\n"
+                    f" {(cache_entry.size/GB):.3f} GB.\n"
                     f"{get_pretty_snapshot_diff(snapshot_before, snapshot_after)}"
                 )
 
     def print_cuda_stats(self) -> None:
         """Log CUDA diagnostics."""
-        vram = "%4.2fG" % (torch.cuda.memory_allocated() / GIG)
-        ram = "%4.2fG" % (self.cache_size() / GIG)
+        vram = "%4.2fG" % (torch.cuda.memory_allocated() / GB)
+        ram = "%4.2fG" % (self.cache_size() / GB)
 
         in_ram_models = 0
         in_vram_models = 0
@@ -353,17 +349,20 @@ class ModelCache(ModelCacheBase[AnyModel]):
                 )
 
     def make_room(self, size: int) -> None:
-        """Make enough room in the cache to accommodate a new model of indicated size."""
-        # calculate how much memory this model will require
-        # multiplier = 2 if self.precision==torch.float32 else 1
+        """Make enough room in the cache to accommodate a new model of indicated size.
+
+        Note: This function deletes all of the cache's internal references to a model in order to free it. If there are
+        external references to the model, there's nothing that the cache can do about it, and those models will not be
+        garbage-collected.
+        """
         bytes_needed = size
-        maximum_size = self.max_cache_size * GIG  # stored in GB, convert to bytes
+        maximum_size = self.max_cache_size * GB  # stored in GB, convert to bytes
         current_size = self.cache_size()
 
         if current_size + bytes_needed > maximum_size:
             self.logger.debug(
-                f"Max cache size exceeded: {(current_size/GIG):.2f}/{self.max_cache_size:.2f} GB, need an additional"
-                f" {(bytes_needed/GIG):.2f} GB"
+                f"Max cache size exceeded: {(current_size/GB):.2f}/{self.max_cache_size:.2f} GB, need an additional"
+                f" {(bytes_needed/GB):.2f} GB"
             )
 
         self.logger.debug(f"Before making_room: cached_models={len(self._cached_models)}")
@@ -380,7 +379,7 @@ class ModelCache(ModelCacheBase[AnyModel]):
 
             if not cache_entry.locked:
                 self.logger.debug(
-                    f"Removing {model_key} from RAM cache to free at least {(size/GIG):.2f} GB (-{(cache_entry.size/GIG):.2f} GB)"
+                    f"Removing {model_key} from RAM cache to free at least {(size/GB):.2f} GB (-{(cache_entry.size/GB):.2f} GB)"
                 )
                 current_size -= cache_entry.size
                 models_cleared += 1

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

@@ -0,0 +1,234 @@
+# Copyright (c) 2024, Brandon W. Rising and the InvokeAI Development Team
+"""Class for Flux model loading in InvokeAI."""
+
+from pathlib import Path
+from typing import Optional
+
+import accelerate
+import torch
+from safetensors.torch import load_file
+from transformers import AutoConfig, AutoModelForTextEncoding, CLIPTextModel, CLIPTokenizer, T5EncoderModel, T5Tokenizer
+
+from invokeai.app.services.config.config_default import get_config
+from invokeai.backend.flux.model import Flux
+from invokeai.backend.flux.modules.autoencoder import AutoEncoder
+from invokeai.backend.flux.util import ae_params, params
+from invokeai.backend.model_manager import (
+    AnyModel,
+    AnyModelConfig,
+    BaseModelType,
+    ModelFormat,
+    ModelType,
+    SubModelType,
+)
+from invokeai.backend.model_manager.config import (
+    CheckpointConfigBase,
+    CLIPEmbedDiffusersConfig,
+    MainBnbQuantized4bCheckpointConfig,
+    MainCheckpointConfig,
+    T5EncoderBnbQuantizedLlmInt8bConfig,
+    T5EncoderConfig,
+    VAECheckpointConfig,
+)
+from invokeai.backend.model_manager.load.load_default import ModelLoader
+from invokeai.backend.model_manager.load.model_loader_registry import ModelLoaderRegistry
+from invokeai.backend.util.silence_warnings import SilenceWarnings
+
+try:
+    from invokeai.backend.quantization.bnb_llm_int8 import quantize_model_llm_int8
+    from invokeai.backend.quantization.bnb_nf4 import quantize_model_nf4
+
+    bnb_available = True
+except ImportError:
+    bnb_available = False
+
+app_config = get_config()
+
+
+@ModelLoaderRegistry.register(base=BaseModelType.Flux, type=ModelType.VAE, format=ModelFormat.Checkpoint)
+class FluxVAELoader(ModelLoader):
+    """Class to load VAE models."""
+
+    def _load_model(
+        self,
+        config: AnyModelConfig,
+        submodel_type: Optional[SubModelType] = None,
+    ) -> AnyModel:
+        if not isinstance(config, VAECheckpointConfig):
+            raise ValueError("Only VAECheckpointConfig models are currently supported here.")
+        model_path = Path(config.path)
+
+        with SilenceWarnings():
+            model = AutoEncoder(ae_params[config.config_path])
+            sd = load_file(model_path)
+            model.load_state_dict(sd, assign=True)
+            model.to(dtype=self._torch_dtype)
+
+        return model
+
+
+@ModelLoaderRegistry.register(base=BaseModelType.Any, type=ModelType.CLIPEmbed, format=ModelFormat.Diffusers)
+class ClipCheckpointModel(ModelLoader):
+    """Class to load main models."""
+
+    def _load_model(
+        self,
+        config: AnyModelConfig,
+        submodel_type: Optional[SubModelType] = None,
+    ) -> AnyModel:
+        if not isinstance(config, CLIPEmbedDiffusersConfig):
+            raise ValueError("Only CLIPEmbedDiffusersConfig models are currently supported here.")
+
+        match submodel_type:
+            case SubModelType.Tokenizer:
+                return CLIPTokenizer.from_pretrained(Path(config.path) / "tokenizer")
+            case SubModelType.TextEncoder:
+                return CLIPTextModel.from_pretrained(Path(config.path) / "text_encoder")
+
+        raise ValueError(
+            f"Only Tokenizer and TextEncoder submodels are currently supported. Received: {submodel_type.value if submodel_type else 'None'}"
+        )
+
+
+@ModelLoaderRegistry.register(base=BaseModelType.Any, type=ModelType.T5Encoder, format=ModelFormat.BnbQuantizedLlmInt8b)
+class BnbQuantizedLlmInt8bCheckpointModel(ModelLoader):
+    """Class to load main models."""
+
+    def _load_model(
+        self,
+        config: AnyModelConfig,
+        submodel_type: Optional[SubModelType] = None,
+    ) -> AnyModel:
+        if not isinstance(config, T5EncoderBnbQuantizedLlmInt8bConfig):
+            raise ValueError("Only T5EncoderBnbQuantizedLlmInt8bConfig models are currently supported here.")
+        if not bnb_available:
+            raise ImportError(
+                "The bnb modules are not available. Please install bitsandbytes if available on your platform."
+            )
+        match submodel_type:
+            case SubModelType.Tokenizer2:
+                return T5Tokenizer.from_pretrained(Path(config.path) / "tokenizer_2", max_length=512)
+            case SubModelType.TextEncoder2:
+                te2_model_path = Path(config.path) / "text_encoder_2"
+                model_config = AutoConfig.from_pretrained(te2_model_path)
+                with accelerate.init_empty_weights():
+                    model = AutoModelForTextEncoding.from_config(model_config)
+                    model = quantize_model_llm_int8(model, modules_to_not_convert=set())
+
+                state_dict_path = te2_model_path / "bnb_llm_int8_model.safetensors"
+                state_dict = load_file(state_dict_path)
+                self._load_state_dict_into_t5(model, state_dict)
+
+                return model
+
+        raise ValueError(
+            f"Only Tokenizer and TextEncoder submodels are currently supported. Received: {submodel_type.value if submodel_type else 'None'}"
+        )
+
+    @classmethod
+    def _load_state_dict_into_t5(cls, model: T5EncoderModel, state_dict: dict[str, torch.Tensor]):
+        # There is a shared reference to a single weight tensor in the model.
+        # Both "encoder.embed_tokens.weight" and "shared.weight" refer to the same tensor, so only the latter should
+        # be present in the state_dict.
+        missing_keys, unexpected_keys = model.load_state_dict(state_dict, strict=False, assign=True)
+        assert len(unexpected_keys) == 0
+        assert set(missing_keys) == {"encoder.embed_tokens.weight"}
+        # Assert that the layers we expect to be shared are actually shared.
+        assert model.encoder.embed_tokens.weight is model.shared.weight
+
+
+@ModelLoaderRegistry.register(base=BaseModelType.Any, type=ModelType.T5Encoder, format=ModelFormat.T5Encoder)
+class T5EncoderCheckpointModel(ModelLoader):
+    """Class to load main models."""
+
+    def _load_model(
+        self,
+        config: AnyModelConfig,
+        submodel_type: Optional[SubModelType] = None,
+    ) -> AnyModel:
+        if not isinstance(config, T5EncoderConfig):
+            raise ValueError("Only T5EncoderConfig models are currently supported here.")
+
+        match submodel_type:
+            case SubModelType.Tokenizer2:
+                return T5Tokenizer.from_pretrained(Path(config.path) / "tokenizer_2", max_length=512)
+            case SubModelType.TextEncoder2:
+                return T5EncoderModel.from_pretrained(Path(config.path) / "text_encoder_2")
+
+        raise ValueError(
+            f"Only Tokenizer and TextEncoder submodels are currently supported. Received: {submodel_type.value if submodel_type else 'None'}"
+        )
+
+
+@ModelLoaderRegistry.register(base=BaseModelType.Flux, type=ModelType.Main, format=ModelFormat.Checkpoint)
+class FluxCheckpointModel(ModelLoader):
+    """Class to load main models."""
+
+    def _load_model(
+        self,
+        config: AnyModelConfig,
+        submodel_type: Optional[SubModelType] = None,
+    ) -> AnyModel:
+        if not isinstance(config, CheckpointConfigBase):
+            raise ValueError("Only CheckpointConfigBase models are currently supported here.")
+
+        match submodel_type:
+            case SubModelType.Transformer:
+                return self._load_from_singlefile(config)
+
+        raise ValueError(
+            f"Only Transformer submodels are currently supported. Received: {submodel_type.value if submodel_type else 'None'}"
+        )
+
+    def _load_from_singlefile(
+        self,
+        config: AnyModelConfig,
+    ) -> AnyModel:
+        assert isinstance(config, MainCheckpointConfig)
+        model_path = Path(config.path)
+
+        with SilenceWarnings():
+            model = Flux(params[config.config_path])
+            sd = load_file(model_path)
+            model.load_state_dict(sd, assign=True)
+        return model
+
+
+@ModelLoaderRegistry.register(base=BaseModelType.Flux, type=ModelType.Main, format=ModelFormat.BnbQuantizednf4b)
+class FluxBnbQuantizednf4bCheckpointModel(ModelLoader):
+    """Class to load main models."""
+
+    def _load_model(
+        self,
+        config: AnyModelConfig,
+        submodel_type: Optional[SubModelType] = None,
+    ) -> AnyModel:
+        if not isinstance(config, CheckpointConfigBase):
+            raise ValueError("Only CheckpointConfigBase models are currently supported here.")
+
+        match submodel_type:
+            case SubModelType.Transformer:
+                return self._load_from_singlefile(config)
+
+        raise ValueError(
+            f"Only Transformer submodels are currently supported. Received: {submodel_type.value if submodel_type else 'None'}"
+        )
+
+    def _load_from_singlefile(
+        self,
+        config: AnyModelConfig,
+    ) -> AnyModel:
+        assert isinstance(config, MainBnbQuantized4bCheckpointConfig)
+        if not bnb_available:
+            raise ImportError(
+                "The bnb modules are not available. Please install bitsandbytes if available on your platform."
+            )
+        model_path = Path(config.path)
+
+        with SilenceWarnings():
+            with accelerate.init_empty_weights():
+                model = Flux(params[config.config_path])
+                model = quantize_model_nf4(model, modules_to_not_convert=set(), compute_dtype=torch.bfloat16)
+            sd = load_file(model_path)
+            model.load_state_dict(sd, assign=True)
+        return model

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

@@ -78,7 +78,12 @@ class GenericDiffusersLoader(ModelLoader):
 
     # TO DO: Add exception handling
     def _hf_definition_to_type(self, module: str, class_name: str) -> ModelMixin:  # fix with correct type
-        if module in ["diffusers", "transformers"]:
+        if module in [
+            "diffusers",
+            "transformers",
+            "invokeai.backend.quantization.fast_quantized_transformers_model",
+            "invokeai.backend.quantization.fast_quantized_diffusion_model",
+        ]:
             res_type = sys.modules[module]
         else:
             res_type = sys.modules["diffusers"].pipelines

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

@@ -36,8 +36,18 @@ VARIANT_TO_IN_CHANNEL_MAP = {
 }
 
 
-@ModelLoaderRegistry.register(base=BaseModelType.Any, type=ModelType.Main, format=ModelFormat.Diffusers)
-@ModelLoaderRegistry.register(base=BaseModelType.Any, type=ModelType.Main, format=ModelFormat.Checkpoint)
+@ModelLoaderRegistry.register(base=BaseModelType.StableDiffusion1, type=ModelType.Main, format=ModelFormat.Diffusers)
+@ModelLoaderRegistry.register(base=BaseModelType.StableDiffusion2, type=ModelType.Main, format=ModelFormat.Diffusers)
+@ModelLoaderRegistry.register(base=BaseModelType.StableDiffusionXL, type=ModelType.Main, format=ModelFormat.Diffusers)
+@ModelLoaderRegistry.register(
+    base=BaseModelType.StableDiffusionXLRefiner, type=ModelType.Main, format=ModelFormat.Diffusers
+)
+@ModelLoaderRegistry.register(base=BaseModelType.StableDiffusion1, type=ModelType.Main, format=ModelFormat.Checkpoint)
+@ModelLoaderRegistry.register(base=BaseModelType.StableDiffusion2, type=ModelType.Main, format=ModelFormat.Checkpoint)
+@ModelLoaderRegistry.register(base=BaseModelType.StableDiffusionXL, type=ModelType.Main, format=ModelFormat.Checkpoint)
+@ModelLoaderRegistry.register(
+    base=BaseModelType.StableDiffusionXLRefiner, type=ModelType.Main, format=ModelFormat.Checkpoint
+)
 class StableDiffusionDiffusersModel(GenericDiffusersLoader):
     """Class to load main models."""
 

invokeai/backend/model_manager/load/model_util.py

@@ -9,7 +9,7 @@ from typing import Optional
 import torch
 from diffusers.pipelines.pipeline_utils import DiffusionPipeline
 from diffusers.schedulers.scheduling_utils import SchedulerMixin
-from transformers import CLIPTokenizer
+from transformers import CLIPTokenizer, T5Tokenizer, T5TokenizerFast
 
 from invokeai.backend.image_util.depth_anything.depth_anything_pipeline import DepthAnythingPipeline
 from invokeai.backend.image_util.grounding_dino.grounding_dino_pipeline import GroundingDinoPipeline
@@ -50,6 +50,17 @@ def calc_model_size_by_data(logger: logging.Logger, model: AnyModel) -> int:
         ),
     ):
         return model.calc_size()
+    elif isinstance(
+        model,
+        (
+            T5TokenizerFast,
+            T5Tokenizer,
+        ),
+    ):
+        # HACK(ryand): len(model) just returns the vocabulary size, so this is blatantly wrong. It should be small
+        # relative to the text encoder that it's used with, so shouldn't matter too much, but we should fix this at some
+        # point.
+        return len(model)
     else:
         # TODO(ryand): Promote this from a log to an exception once we are confident that we are handling all of the
         # supported model types.

invokeai/backend/model_manager/probe.py

@@ -95,6 +95,7 @@ class ModelProbe(object):
     }
 
     CLASS2TYPE = {
+        "FluxPipeline": ModelType.Main,
         "StableDiffusionPipeline": ModelType.Main,
         "StableDiffusionInpaintPipeline": ModelType.Main,
         "StableDiffusionXLPipeline": ModelType.Main,
@@ -106,6 +107,7 @@ class ModelProbe(object):
         "ControlNetModel": ModelType.ControlNet,
         "CLIPVisionModelWithProjection": ModelType.CLIPVision,
         "T2IAdapter": ModelType.T2IAdapter,
+        "CLIPModel": ModelType.CLIPEmbed,
     }
 
     @classmethod
@@ -161,7 +163,7 @@ class ModelProbe(object):
         fields["description"] = (
             fields.get("description") or f"{fields['base'].value} {model_type.value} model {fields['name']}"
         )
-        fields["format"] = fields.get("format") or probe.get_format()
+        fields["format"] = ModelFormat(fields.get("format")) if "format" in fields else probe.get_format()
         fields["hash"] = fields.get("hash") or ModelHash(algorithm=hash_algo).hash(model_path)
 
         fields["default_settings"] = fields.get("default_settings")
@@ -176,10 +178,10 @@ class ModelProbe(object):
             fields["repo_variant"] = fields.get("repo_variant") or probe.get_repo_variant()
 
         # additional fields needed for main and controlnet models
-        if (
-            fields["type"] in [ModelType.Main, ModelType.ControlNet, ModelType.VAE]
-            and fields["format"] is ModelFormat.Checkpoint
-        ):
+        if fields["type"] in [ModelType.Main, ModelType.ControlNet, ModelType.VAE] and fields["format"] in [
+            ModelFormat.Checkpoint,
+            ModelFormat.BnbQuantizednf4b,
+        ]:
             ckpt_config_path = cls._get_checkpoint_config_path(
                 model_path,
                 model_type=fields["type"],
@@ -222,7 +224,8 @@ class ModelProbe(object):
         ckpt = ckpt.get("state_dict", ckpt)
 
         for key in [str(k) for k in ckpt.keys()]:
-            if key.startswith(("cond_stage_model.", "first_stage_model.", "model.diffusion_model.")):
+            if key.startswith(("cond_stage_model.", "first_stage_model.", "model.diffusion_model.", "double_blocks.")):
+                # Keys starting with double_blocks are associated with Flux models
                 return ModelType.Main
             elif key.startswith(("encoder.conv_in", "decoder.conv_in")):
                 return ModelType.VAE
@@ -321,10 +324,27 @@ class ModelProbe(object):
             return possible_conf.absolute()
 
         if model_type is ModelType.Main:
-            config_file = LEGACY_CONFIGS[base_type][variant_type]
-            if isinstance(config_file, dict):  # need another tier for sd-2.x models
-                config_file = config_file[prediction_type]
-            config_file = f"stable-diffusion/{config_file}"
+            if base_type == BaseModelType.Flux:
+                # TODO: Decide between dev/schnell
+                checkpoint = ModelProbe._scan_and_load_checkpoint(model_path)
+                state_dict = checkpoint.get("state_dict") or checkpoint
+                if "guidance_in.out_layer.weight" in state_dict:
+                    # For flux, this is a key in invokeai.backend.flux.util.params
+                    #   Due to model type and format being the descriminator for model configs this
+                    #   is used rather than attempting to support flux with separate model types and format
+                    #   If changed in the future, please fix me
+                    config_file = "flux-dev"
+                else:
+                    # For flux, this is a key in invokeai.backend.flux.util.params
+                    #   Due to model type and format being the descriminator for model configs this
+                    #   is used rather than attempting to support flux with separate model types and format
+                    #   If changed in the future, please fix me
+                    config_file = "flux-schnell"
+            else:
+                config_file = LEGACY_CONFIGS[base_type][variant_type]
+                if isinstance(config_file, dict):  # need another tier for sd-2.x models
+                    config_file = config_file[prediction_type]
+                config_file = f"stable-diffusion/{config_file}"
         elif model_type is ModelType.ControlNet:
             config_file = (
                 "controlnet/cldm_v15.yaml"
@@ -333,7 +353,13 @@ class ModelProbe(object):
             )
         elif model_type is ModelType.VAE:
             config_file = (
-                "stable-diffusion/v1-inference.yaml"
+                # For flux, this is a key in invokeai.backend.flux.util.ae_params
+                #   Due to model type and format being the descriminator for model configs this
+                #   is used rather than attempting to support flux with separate model types and format
+                #   If changed in the future, please fix me
+                "flux"
+                if base_type is BaseModelType.Flux
+                else "stable-diffusion/v1-inference.yaml"
                 if base_type is BaseModelType.StableDiffusion1
                 else "stable-diffusion/sd_xl_base.yaml"
                 if base_type is BaseModelType.StableDiffusionXL
@@ -416,11 +442,15 @@ class CheckpointProbeBase(ProbeBase):
         self.checkpoint = ModelProbe._scan_and_load_checkpoint(model_path)
 
     def get_format(self) -> ModelFormat:
+        state_dict = self.checkpoint.get("state_dict") or self.checkpoint
+        if "double_blocks.0.img_attn.proj.weight.quant_state.bitsandbytes__nf4" in state_dict:
+            return ModelFormat.BnbQuantizednf4b
         return ModelFormat("checkpoint")
 
     def get_variant_type(self) -> ModelVariantType:
         model_type = ModelProbe.get_model_type_from_checkpoint(self.model_path, self.checkpoint)
-        if model_type != ModelType.Main:
+        base_type = self.get_base_type()
+        if model_type != ModelType.Main or base_type == BaseModelType.Flux:
             return ModelVariantType.Normal
         state_dict = self.checkpoint.get("state_dict") or self.checkpoint
         in_channels = state_dict["model.diffusion_model.input_blocks.0.0.weight"].shape[1]
@@ -440,6 +470,8 @@ class PipelineCheckpointProbe(CheckpointProbeBase):
     def get_base_type(self) -> BaseModelType:
         checkpoint = self.checkpoint
         state_dict = self.checkpoint.get("state_dict") or checkpoint
+        if "double_blocks.0.img_attn.norm.key_norm.scale" in state_dict:
+            return BaseModelType.Flux
         key_name = "model.diffusion_model.input_blocks.2.1.transformer_blocks.0.attn2.to_k.weight"
         if key_name in state_dict and state_dict[key_name].shape[-1] == 768:
             return BaseModelType.StableDiffusion1
@@ -482,6 +514,7 @@ class VaeCheckpointProbe(CheckpointProbeBase):
             (r"xl", BaseModelType.StableDiffusionXL),
             (r"sd2", BaseModelType.StableDiffusion2),
             (r"vae", BaseModelType.StableDiffusion1),
+            (r"FLUX.1-schnell_ae", BaseModelType.Flux),
         ]:
             if re.search(regexp, self.model_path.name, re.IGNORECASE):
                 return basetype
@@ -713,6 +746,11 @@ class TextualInversionFolderProbe(FolderProbeBase):
         return TextualInversionCheckpointProbe(path).get_base_type()
 
 
+class T5EncoderFolderProbe(FolderProbeBase):
+    def get_format(self) -> ModelFormat:
+        return ModelFormat.T5Encoder
+
+
 class ONNXFolderProbe(PipelineFolderProbe):
     def get_base_type(self) -> BaseModelType:
         # Due to the way the installer is set up, the configuration file for safetensors
@@ -805,6 +843,11 @@ class CLIPVisionFolderProbe(FolderProbeBase):
         return BaseModelType.Any
 
 
+class CLIPEmbedFolderProbe(FolderProbeBase):
+    def get_base_type(self) -> BaseModelType:
+        return BaseModelType.Any
+
+
 class SpandrelImageToImageFolderProbe(FolderProbeBase):
     def get_base_type(self) -> BaseModelType:
         raise NotImplementedError()
@@ -835,8 +878,10 @@ ModelProbe.register_probe("diffusers", ModelType.Main, PipelineFolderProbe)
 ModelProbe.register_probe("diffusers", ModelType.VAE, VaeFolderProbe)
 ModelProbe.register_probe("diffusers", ModelType.LoRA, LoRAFolderProbe)
 ModelProbe.register_probe("diffusers", ModelType.TextualInversion, TextualInversionFolderProbe)
+ModelProbe.register_probe("diffusers", ModelType.T5Encoder, T5EncoderFolderProbe)
 ModelProbe.register_probe("diffusers", ModelType.ControlNet, ControlNetFolderProbe)
 ModelProbe.register_probe("diffusers", ModelType.IPAdapter, IPAdapterFolderProbe)
+ModelProbe.register_probe("diffusers", ModelType.CLIPEmbed, CLIPEmbedFolderProbe)
 ModelProbe.register_probe("diffusers", ModelType.CLIPVision, CLIPVisionFolderProbe)
 ModelProbe.register_probe("diffusers", ModelType.T2IAdapter, T2IAdapterFolderProbe)
 ModelProbe.register_probe("diffusers", ModelType.SpandrelImageToImage, SpandrelImageToImageFolderProbe)

invokeai/backend/model_manager/starter_models.py

@@ -2,7 +2,7 @@ from typing import Optional
 
 from pydantic import BaseModel
 
-from invokeai.backend.model_manager.config import BaseModelType, ModelType
+from invokeai.backend.model_manager.config import BaseModelType, ModelFormat, ModelType
 
 
 class StarterModelWithoutDependencies(BaseModel):
@@ -11,6 +11,7 @@ class StarterModelWithoutDependencies(BaseModel):
     name: str
     base: BaseModelType
     type: ModelType
+    format: Optional[ModelFormat] = None
     is_installed: bool = False
 
 
@@ -51,10 +52,76 @@ cyberrealistic_negative = StarterModel(
     type=ModelType.TextualInversion,
 )
 
+t5_base_encoder = StarterModel(
+    name="t5_base_encoder",
+    base=BaseModelType.Any,
+    source="InvokeAI/t5-v1_1-xxl::bfloat16",
+    description="T5-XXL text encoder (used in FLUX pipelines). ~8GB",
+    type=ModelType.T5Encoder,
+)
+
+t5_8b_quantized_encoder = StarterModel(
+    name="t5_bnb_int8_quantized_encoder",
+    base=BaseModelType.Any,
+    source="InvokeAI/t5-v1_1-xxl::bnb_llm_int8",
+    description="T5-XXL text encoder with bitsandbytes LLM.int8() quantization (used in FLUX pipelines). ~5GB",
+    type=ModelType.T5Encoder,
+    format=ModelFormat.BnbQuantizedLlmInt8b,
+)
+
+clip_l_encoder = StarterModel(
+    name="clip-vit-large-patch14",
+    base=BaseModelType.Any,
+    source="InvokeAI/clip-vit-large-patch14-text-encoder::bfloat16",
+    description="CLIP-L text encoder (used in FLUX pipelines). ~250MB",
+    type=ModelType.CLIPEmbed,
+)
+
+flux_vae = StarterModel(
+    name="FLUX.1-schnell_ae",
+    base=BaseModelType.Flux,
+    source="black-forest-labs/FLUX.1-schnell::ae.safetensors",
+    description="FLUX VAE compatible with both schnell and dev variants.",
+    type=ModelType.VAE,
+)
+
+
 # List of starter models, displayed on the frontend.
 # The order/sort of this list is not changed by the frontend - set it how you want it here.
 STARTER_MODELS: list[StarterModel] = [
     # region: Main
+    StarterModel(
+        name="FLUX Schnell (Quantized)",
+        base=BaseModelType.Flux,
+        source="InvokeAI/flux_schnell::transformer/bnb_nf4/flux1-schnell-bnb_nf4.safetensors",
+        description="FLUX schnell transformer quantized to bitsandbytes NF4 format. Total size with dependencies: ~12GB",
+        type=ModelType.Main,
+        dependencies=[t5_8b_quantized_encoder, flux_vae, clip_l_encoder],
+    ),
+    StarterModel(
+        name="FLUX Dev (Quantized)",
+        base=BaseModelType.Flux,
+        source="InvokeAI/flux_dev::transformer/bnb_nf4/flux1-dev-bnb_nf4.safetensors",
+        description="FLUX dev transformer quantized to bitsandbytes NF4 format. Total size with dependencies: ~12GB",
+        type=ModelType.Main,
+        dependencies=[t5_8b_quantized_encoder, flux_vae, clip_l_encoder],
+    ),
+    StarterModel(
+        name="FLUX Schnell",
+        base=BaseModelType.Flux,
+        source="InvokeAI/flux_schnell::transformer/base/flux1-schnell.safetensors",
+        description="FLUX schnell transformer in bfloat16. Total size with dependencies: ~33GB",
+        type=ModelType.Main,
+        dependencies=[t5_base_encoder, flux_vae, clip_l_encoder],
+    ),
+    StarterModel(
+        name="FLUX Dev",
+        base=BaseModelType.Flux,
+        source="InvokeAI/flux_dev::transformer/base/flux1-dev.safetensors",
+        description="FLUX dev transformer in bfloat16. Total size with dependencies: ~33GB",
+        type=ModelType.Main,
+        dependencies=[t5_base_encoder, flux_vae, clip_l_encoder],
+    ),
     StarterModel(
         name="CyberRealistic v4.1",
         base=BaseModelType.StableDiffusion1,
@@ -125,6 +192,7 @@ STARTER_MODELS: list[StarterModel] = [
     # endregion
     # region VAE
     sdxl_fp16_vae_fix,
+    flux_vae,
     # endregion
     # region LoRA
     StarterModel(
@@ -450,6 +518,11 @@ STARTER_MODELS: list[StarterModel] = [
         type=ModelType.SpandrelImageToImage,
     ),
     # endregion
+    # region TextEncoders
+    t5_base_encoder,
+    t5_8b_quantized_encoder,
+    clip_l_encoder,
+    # endregion
 ]
 
 assert len(STARTER_MODELS) == len({m.source for m in STARTER_MODELS}), "Duplicate starter models"

invokeai/backend/model_manager/util/select_hf_files.py

@@ -54,6 +54,7 @@ def filter_files(
                 "lora_weights.safetensors",
                 "weights.pb",
                 "onnx_data",
+                "spiece.model",  # Added for `black-forest-labs/FLUX.1-schnell`.
             )
         ):
             paths.append(file)
@@ -62,13 +63,13 @@ def filter_files(
         # downloading random checkpoints that might also be in the repo. However there is no guarantee
         # that a checkpoint doesn't contain "model" in its name, and no guarantee that future diffusers models
         # will adhere to this naming convention, so this is an area to be careful of.
-        elif re.search(r"model(\.[^.]+)?\.(safetensors|bin|onnx|xml|pth|pt|ckpt|msgpack)$", file.name):
+        elif re.search(r"model.*\.(safetensors|bin|onnx|xml|pth|pt|ckpt|msgpack)$", file.name):
             paths.append(file)
 
     # limit search to subfolder if requested
     if subfolder:
         subfolder = root / subfolder
-        paths = [x for x in paths if x.parent == Path(subfolder)]
+        paths = [x for x in paths if Path(subfolder) in x.parents]
 
     # _filter_by_variant uniquifies the paths and returns a set
     return sorted(_filter_by_variant(paths, variant))
@@ -97,7 +98,9 @@ def _filter_by_variant(files: List[Path], variant: ModelRepoVariant) -> Set[Path
             if variant == ModelRepoVariant.Flax:
                 result.add(path)
 
-        elif path.suffix in [".json", ".txt"]:
+        # Note: '.model' was added to support:
+        # https://huggingface.co/black-forest-labs/FLUX.1-schnell/blob/768d12a373ed5cc9ef9a9dea7504dc09fcc14842/tokenizer_2/spiece.model
+        elif path.suffix in [".json", ".txt", ".model"]:
             result.add(path)
 
         elif variant in [
@@ -140,6 +143,23 @@ def _filter_by_variant(files: List[Path], variant: ModelRepoVariant) -> Set[Path
             continue
 
     for candidate_list in subfolder_weights.values():
+        # Check if at least one of the files has the explicit fp16 variant.
+        at_least_one_fp16 = False
+        for candidate in candidate_list:
+            if len(candidate.path.suffixes) == 2 and candidate.path.suffixes[0] == ".fp16":
+                at_least_one_fp16 = True
+                break
+
+        if not at_least_one_fp16:
+            # If none of the candidates in this candidate_list have the explicit fp16 variant label, then this
+            # candidate_list probably doesn't adhere to the variant naming convention that we expected. In this case,
+            # we'll simply keep all the candidates. An example of a model that hits this case is
+            # `black-forest-labs/FLUX.1-schnell` (as of commit 012d2fd).
+            for candidate in candidate_list:
+                result.add(candidate.path)
+
+        # The candidate_list seems to have the expected variant naming convention. We'll select the highest scoring
+        # candidate.
         highest_score_candidate = max(candidate_list, key=lambda candidate: candidate.score)
         if highest_score_candidate:
             result.add(highest_score_candidate.path)

invokeai/backend/quantization/bnb_llm_int8.py

@@ -0,0 +1,135 @@
+import bitsandbytes as bnb
+import torch
+
+# This file contains utils for working with models that use bitsandbytes LLM.int8() quantization.
+# The utils in this file are partially inspired by:
+# https://github.com/Lightning-AI/pytorch-lightning/blob/1551a16b94f5234a4a78801098f64d0732ef5cb5/src/lightning/fabric/plugins/precision/bitsandbytes.py
+
+
+# NOTE(ryand): All of the custom state_dict manipulation logic in this file is pretty hacky. This could be made much
+# cleaner by re-implementing bnb.nn.Linear8bitLt with proper use of buffers and less magic. But, for now, we try to
+# stick close to the bitsandbytes classes to make interoperability easier with other models that might use bitsandbytes.
+
+
+class InvokeInt8Params(bnb.nn.Int8Params):
+    """We override cuda() to avoid re-quantizing the weights in the following cases:
+    - We loaded quantized weights from a state_dict on the cpu, and then moved the model to the gpu.
+    - We are moving the model back-and-forth between the cpu and gpu.
+    """
+
+    def cuda(self, device):
+        if self.has_fp16_weights:
+            return super().cuda(device)
+        elif self.CB is not None and self.SCB is not None:
+            self.data = self.data.cuda()
+            self.CB = self.data
+            self.SCB = self.SCB.cuda()
+        else:
+            # we store the 8-bit rows-major weight
+            # we convert this weight to the turning/ampere weight during the first inference pass
+            B = self.data.contiguous().half().cuda(device)
+            CB, CBt, SCB, SCBt, coo_tensorB = bnb.functional.double_quant(B)
+            del CBt
+            del SCBt
+            self.data = CB
+            self.CB = CB
+            self.SCB = SCB
+
+        return self
+
+
+class InvokeLinear8bitLt(bnb.nn.Linear8bitLt):
+    def _load_from_state_dict(
+        self,
+        state_dict: dict[str, torch.Tensor],
+        prefix: str,
+        local_metadata,
+        strict,
+        missing_keys,
+        unexpected_keys,
+        error_msgs,
+    ):
+        weight = state_dict.pop(prefix + "weight")
+        bias = state_dict.pop(prefix + "bias", None)
+
+        # See `bnb.nn.Linear8bitLt._save_to_state_dict()` for the serialization logic of SCB and weight_format.
+        scb = state_dict.pop(prefix + "SCB", None)
+
+        # Currently, we only support weight_format=0.
+        weight_format = state_dict.pop(prefix + "weight_format", None)
+        assert weight_format == 0
+
+        # TODO(ryand): Technically, we should be using `strict`, `missing_keys`, `unexpected_keys`, and `error_msgs`
+        # rather than raising an exception to correctly implement this API.
+        assert len(state_dict) == 0
+
+        if scb is not None:
+            # We are loading a pre-quantized state dict.
+            self.weight = InvokeInt8Params(
+                data=weight,
+                requires_grad=self.weight.requires_grad,
+                has_fp16_weights=False,
+                # Note: After quantization, CB is the same as weight.
+                CB=weight,
+                SCB=scb,
+            )
+            self.bias = bias if bias is None else torch.nn.Parameter(bias)
+        else:
+            # We are loading a non-quantized state dict.
+
+            # We could simply call the `super()._load_from_state_dict()` method here, but then we wouldn't be able to
+            # load from a state_dict into a model on the "meta" device. Attempting to load into a model on the "meta"
+            # device requires setting `assign=True`, doing this with the default `super()._load_from_state_dict()`
+            # implementation causes `Params4Bit` to be replaced by a `torch.nn.Parameter`. By initializing a new
+            # `Params4bit` object, we work around this issue. It's a bit hacky, but it gets the job done.
+            self.weight = InvokeInt8Params(
+                data=weight,
+                requires_grad=self.weight.requires_grad,
+                has_fp16_weights=False,
+                CB=None,
+                SCB=None,
+            )
+            self.bias = bias if bias is None else torch.nn.Parameter(bias)
+
+        # Reset the state. The persisted fields are based on the initialization behaviour in
+        # `bnb.nn.Linear8bitLt.__init__()`.
+        new_state = bnb.MatmulLtState()
+        new_state.threshold = self.state.threshold
+        new_state.has_fp16_weights = False
+        new_state.use_pool = self.state.use_pool
+        self.state = new_state
+
+
+def _convert_linear_layers_to_llm_8bit(
+    module: torch.nn.Module, ignore_modules: set[str], outlier_threshold: float, prefix: str = ""
+) -> None:
+    """Convert all linear layers in the module to bnb.nn.Linear8bitLt layers."""
+    for name, child in module.named_children():
+        fullname = f"{prefix}.{name}" if prefix else name
+        if isinstance(child, torch.nn.Linear) and not any(fullname.startswith(s) for s in ignore_modules):
+            has_bias = child.bias is not None
+            replacement = InvokeLinear8bitLt(
+                child.in_features,
+                child.out_features,
+                bias=has_bias,
+                has_fp16_weights=False,
+                threshold=outlier_threshold,
+            )
+            replacement.weight.data = child.weight.data
+            if has_bias:
+                replacement.bias.data = child.bias.data
+            replacement.requires_grad_(False)
+            module.__setattr__(name, replacement)
+        else:
+            _convert_linear_layers_to_llm_8bit(
+                child, ignore_modules, outlier_threshold=outlier_threshold, prefix=fullname
+            )
+
+
+def quantize_model_llm_int8(model: torch.nn.Module, modules_to_not_convert: set[str], outlier_threshold: float = 6.0):
+    """Apply bitsandbytes LLM.8bit() quantization to the model."""
+    _convert_linear_layers_to_llm_8bit(
+        module=model, ignore_modules=modules_to_not_convert, outlier_threshold=outlier_threshold
+    )
+
+    return model

invokeai/backend/quantization/bnb_nf4.py

@@ -0,0 +1,156 @@
+import bitsandbytes as bnb
+import torch
+
+# This file contains utils for working with models that use bitsandbytes NF4 quantization.
+# The utils in this file are partially inspired by:
+# https://github.com/Lightning-AI/pytorch-lightning/blob/1551a16b94f5234a4a78801098f64d0732ef5cb5/src/lightning/fabric/plugins/precision/bitsandbytes.py
+
+# NOTE(ryand): All of the custom state_dict manipulation logic in this file is pretty hacky. This could be made much
+# cleaner by re-implementing bnb.nn.LinearNF4 with proper use of buffers and less magic. But, for now, we try to stick
+# close to the bitsandbytes classes to make interoperability easier with other models that might use bitsandbytes.
+
+
+class InvokeLinearNF4(bnb.nn.LinearNF4):
+    """A class that extends `bnb.nn.LinearNF4` to add the following functionality:
+    - Ability to load Linear NF4 layers from a pre-quantized state_dict.
+    - Ability to load Linear NF4 layers from a state_dict when the model is on the "meta" device.
+    """
+
+    def _load_from_state_dict(
+        self,
+        state_dict: dict[str, torch.Tensor],
+        prefix: str,
+        local_metadata,
+        strict,
+        missing_keys,
+        unexpected_keys,
+        error_msgs,
+    ):
+        """This method is based on the logic in the bitsandbytes serialization unit tests for `Linear4bit`:
+        https://github.com/bitsandbytes-foundation/bitsandbytes/blob/6d714a5cce3db5bd7f577bc447becc7a92d5ccc7/tests/test_linear4bit.py#L52-L71
+        """
+        weight = state_dict.pop(prefix + "weight")
+        bias = state_dict.pop(prefix + "bias", None)
+        # We expect the remaining keys to be quant_state keys.
+        quant_state_sd = state_dict
+
+        # During serialization, the quant_state is stored as subkeys of "weight." (See
+        # `bnb.nn.LinearNF4._save_to_state_dict()`). We validate that they at least have the correct prefix.
+        # TODO(ryand): Technically, we should be using `strict`, `missing_keys`, `unexpected_keys`, and `error_msgs`
+        # rather than raising an exception to correctly implement this API.
+        assert all(k.startswith(prefix + "weight.") for k in quant_state_sd.keys())
+
+        if len(quant_state_sd) > 0:
+            # We are loading a pre-quantized state dict.
+            self.weight = bnb.nn.Params4bit.from_prequantized(
+                data=weight, quantized_stats=quant_state_sd, device=weight.device
+            )
+            self.bias = bias if bias is None else torch.nn.Parameter(bias, requires_grad=False)
+        else:
+            # We are loading a non-quantized state dict.
+
+            # We could simply call the `super()._load_from_state_dict()` method here, but then we wouldn't be able to
+            # load from a state_dict into a model on the "meta" device. Attempting to load into a model on the "meta"
+            # device requires setting `assign=True`, doing this with the default `super()._load_from_state_dict()`
+            # implementation causes `Params4Bit` to be replaced by a `torch.nn.Parameter`. By initializing a new
+            # `Params4bit` object, we work around this issue. It's a bit hacky, but it gets the job done.
+            self.weight = bnb.nn.Params4bit(
+                data=weight,
+                requires_grad=self.weight.requires_grad,
+                compress_statistics=self.weight.compress_statistics,
+                quant_type=self.weight.quant_type,
+                quant_storage=self.weight.quant_storage,
+                module=self,
+            )
+            self.bias = bias if bias is None else torch.nn.Parameter(bias)
+
+
+def _replace_param(
+    param: torch.nn.Parameter | bnb.nn.Params4bit,
+    data: torch.Tensor,
+) -> torch.nn.Parameter:
+    """A helper function to replace the data of a model parameter with new data in a way that allows replacing params on
+    the "meta" device.
+
+    Supports both `torch.nn.Parameter` and `bnb.nn.Params4bit` parameters.
+    """
+    if param.device.type == "meta":
+        # Doing `param.data = data` raises a RuntimeError if param.data was on the "meta" device, so we need to
+        # re-create the param instead of overwriting the data.
+        if isinstance(param, bnb.nn.Params4bit):
+            return bnb.nn.Params4bit(
+                data,
+                requires_grad=data.requires_grad,
+                quant_state=param.quant_state,
+                compress_statistics=param.compress_statistics,
+                quant_type=param.quant_type,
+            )
+        return torch.nn.Parameter(data, requires_grad=data.requires_grad)
+
+    param.data = data
+    return param
+
+
+def _convert_linear_layers_to_nf4(
+    module: torch.nn.Module,
+    ignore_modules: set[str],
+    compute_dtype: torch.dtype,
+    compress_statistics: bool = False,
+    prefix: str = "",
+) -> None:
+    """Convert all linear layers in the model to NF4 quantized linear layers.
+
+    Args:
+        module: All linear layers in this module will be converted.
+        ignore_modules: A set of module prefixes to ignore when converting linear layers.
+        compute_dtype: The dtype to use for computation in the quantized linear layers.
+        compress_statistics: Whether to enable nested quantization (aka double quantization) where the quantization
+           constants from the first quantization are quantized again.
+        prefix: The prefix of the current module in the model. Used to call this function recursively.
+    """
+    for name, child in module.named_children():
+        fullname = f"{prefix}.{name}" if prefix else name
+        if isinstance(child, torch.nn.Linear) and not any(fullname.startswith(s) for s in ignore_modules):
+            has_bias = child.bias is not None
+            replacement = InvokeLinearNF4(
+                child.in_features,
+                child.out_features,
+                bias=has_bias,
+                compute_dtype=compute_dtype,
+                compress_statistics=compress_statistics,
+            )
+            if has_bias:
+                replacement.bias = _replace_param(replacement.bias, child.bias.data)
+            replacement.weight = _replace_param(replacement.weight, child.weight.data)
+            replacement.requires_grad_(False)
+            module.__setattr__(name, replacement)
+        else:
+            _convert_linear_layers_to_nf4(child, ignore_modules, compute_dtype=compute_dtype, prefix=fullname)
+
+
+def quantize_model_nf4(model: torch.nn.Module, modules_to_not_convert: set[str], compute_dtype: torch.dtype):
+    """Apply bitsandbytes nf4 quantization to the model.
+
+    You likely want to call this function inside a `accelerate.init_empty_weights()` context.
+
+    Example usage:
+    ```
+    # Initialize the model from a config on the meta device.
+    with accelerate.init_empty_weights():
+        model = ModelClass.from_config(...)
+
+    # Add NF4 quantization linear layers to the model - still on the meta device.
+    with accelerate.init_empty_weights():
+        model = quantize_model_nf4(model, modules_to_not_convert=set(), compute_dtype=torch.float16)
+
+    # Load a state_dict into the model. (Could be either a prequantized or non-quantized state_dict.)
+    model.load_state_dict(state_dict, strict=True, assign=True)
+
+    # Move the model to the "cuda" device. If the model was non-quantized, this is where the weight quantization takes
+    # place.
+    model.to("cuda")
+    ```
+    """
+    _convert_linear_layers_to_nf4(module=model, ignore_modules=modules_to_not_convert, compute_dtype=compute_dtype)
+
+    return model

invokeai/backend/quantization/scripts/load_flux_model_bnb_llm_int8.py

@@ -0,0 +1,79 @@
+from pathlib import Path
+
+import accelerate
+from safetensors.torch import load_file, save_file
+
+from invokeai.backend.flux.model import Flux
+from invokeai.backend.flux.util import params
+from invokeai.backend.quantization.bnb_llm_int8 import quantize_model_llm_int8
+from invokeai.backend.quantization.scripts.load_flux_model_bnb_nf4 import log_time
+
+
+def main():
+    """A script for quantizing a FLUX transformer model using the bitsandbytes LLM.int8() quantization method.
+
+    This script is primarily intended for reference. The script params (e.g. the model_path, modules_to_not_convert,
+    etc.) are hardcoded and would need to be modified for other use cases.
+    """
+    # Load the FLUX transformer model onto the meta device.
+    model_path = Path(
+        "/data/invokeai/models/.download_cache/https__huggingface.co_black-forest-labs_flux.1-schnell_resolve_main_flux1-schnell.safetensors/flux1-schnell.safetensors"
+    )
+
+    with log_time("Intialize FLUX transformer on meta device"):
+        # TODO(ryand): Determine if this is a schnell model or a dev model and load the appropriate config.
+        p = params["flux-schnell"]
+
+        # Initialize the model on the "meta" device.
+        with accelerate.init_empty_weights():
+            model = Flux(p)
+
+    # TODO(ryand): We may want to add some modules to not quantize here (e.g. the proj_out layer). See the accelerate
+    # `get_keys_to_not_convert(...)` function for a heuristic to determine which modules to not quantize.
+    modules_to_not_convert: set[str] = set()
+
+    model_int8_path = model_path.parent / "bnb_llm_int8.safetensors"
+    if model_int8_path.exists():
+        # The quantized model already exists, load it and return it.
+        print(f"A pre-quantized model already exists at '{model_int8_path}'. Attempting to load it...")
+
+        # Replace the linear layers with LLM.int8() quantized linear layers (still on the meta device).
+        with log_time("Replace linear layers with LLM.int8() layers"), accelerate.init_empty_weights():
+            model = quantize_model_llm_int8(model, modules_to_not_convert=modules_to_not_convert)
+
+        with log_time("Load state dict into model"):
+            sd = load_file(model_int8_path)
+            model.load_state_dict(sd, strict=True, assign=True)
+
+        with log_time("Move model to cuda"):
+            model = model.to("cuda")
+
+        print(f"Successfully loaded pre-quantized model from '{model_int8_path}'.")
+
+    else:
+        # The quantized model does not exist, quantize the model and save it.
+        print(f"No pre-quantized model found at '{model_int8_path}'. Quantizing the model...")
+
+        with log_time("Replace linear layers with LLM.int8() layers"), accelerate.init_empty_weights():
+            model = quantize_model_llm_int8(model, modules_to_not_convert=modules_to_not_convert)
+
+        with log_time("Load state dict into model"):
+            state_dict = load_file(model_path)
+            # TODO(ryand): Cast the state_dict to the appropriate dtype?
+            model.load_state_dict(state_dict, strict=True, assign=True)
+
+        with log_time("Move model to cuda and quantize"):
+            model = model.to("cuda")
+
+        with log_time("Save quantized model"):
+            model_int8_path.parent.mkdir(parents=True, exist_ok=True)
+            save_file(model.state_dict(), model_int8_path)
+
+        print(f"Successfully quantized and saved model to '{model_int8_path}'.")
+
+    assert isinstance(model, Flux)
+    return model
+
+
+if __name__ == "__main__":
+    main()

invokeai/backend/quantization/scripts/load_flux_model_bnb_nf4.py

@@ -0,0 +1,96 @@
+import time
+from contextlib import contextmanager
+from pathlib import Path
+
+import accelerate
+import torch
+from safetensors.torch import load_file, save_file
+
+from invokeai.backend.flux.model import Flux
+from invokeai.backend.flux.util import params
+from invokeai.backend.quantization.bnb_nf4 import quantize_model_nf4
+
+
+@contextmanager
+def log_time(name: str):
+    """Helper context manager to log the time taken by a block of code."""
+    start = time.time()
+    try:
+        yield None
+    finally:
+        end = time.time()
+        print(f"'{name}' took {end - start:.4f} secs")
+
+
+def main():
+    """A script for quantizing a FLUX transformer model using the bitsandbytes NF4 quantization method.
+
+    This script is primarily intended for reference. The script params (e.g. the model_path, modules_to_not_convert,
+    etc.) are hardcoded and would need to be modified for other use cases.
+    """
+    model_path = Path(
+        "/data/invokeai/models/.download_cache/https__huggingface.co_black-forest-labs_flux.1-schnell_resolve_main_flux1-schnell.safetensors/flux1-schnell.safetensors"
+    )
+
+    # inference_dtype = torch.bfloat16
+    with log_time("Intialize FLUX transformer on meta device"):
+        # TODO(ryand): Determine if this is a schnell model or a dev model and load the appropriate config.
+        p = params["flux-schnell"]
+
+        # Initialize the model on the "meta" device.
+        with accelerate.init_empty_weights():
+            model = Flux(p)
+
+    # TODO(ryand): We may want to add some modules to not quantize here (e.g. the proj_out layer). See the accelerate
+    # `get_keys_to_not_convert(...)` function for a heuristic to determine which modules to not quantize.
+    modules_to_not_convert: set[str] = set()
+
+    model_nf4_path = model_path.parent / "bnb_nf4.safetensors"
+    if model_nf4_path.exists():
+        # The quantized model already exists, load it and return it.
+        print(f"A pre-quantized model already exists at '{model_nf4_path}'. Attempting to load it...")
+
+        # Replace the linear layers with NF4 quantized linear layers (still on the meta device).
+        with log_time("Replace linear layers with NF4 layers"), accelerate.init_empty_weights():
+            model = quantize_model_nf4(
+                model, modules_to_not_convert=modules_to_not_convert, compute_dtype=torch.bfloat16
+            )
+
+        with log_time("Load state dict into model"):
+            state_dict = load_file(model_nf4_path)
+            model.load_state_dict(state_dict, strict=True, assign=True)
+
+        with log_time("Move model to cuda"):
+            model = model.to("cuda")
+
+        print(f"Successfully loaded pre-quantized model from '{model_nf4_path}'.")
+
+    else:
+        # The quantized model does not exist, quantize the model and save it.
+        print(f"No pre-quantized model found at '{model_nf4_path}'. Quantizing the model...")
+
+        with log_time("Replace linear layers with NF4 layers"), accelerate.init_empty_weights():
+            model = quantize_model_nf4(
+                model, modules_to_not_convert=modules_to_not_convert, compute_dtype=torch.bfloat16
+            )
+
+        with log_time("Load state dict into model"):
+            state_dict = load_file(model_path)
+            # TODO(ryand): Cast the state_dict to the appropriate dtype?
+            model.load_state_dict(state_dict, strict=True, assign=True)
+
+        with log_time("Move model to cuda and quantize"):
+            model = model.to("cuda")
+
+        with log_time("Save quantized model"):
+            model_nf4_path.parent.mkdir(parents=True, exist_ok=True)
+            save_file(model.state_dict(), model_nf4_path)
+
+        print(f"Successfully quantized and saved model to '{model_nf4_path}'.")
+
+    assert isinstance(model, Flux)
+    return model
+
+
+if __name__ == "__main__":
+    main()

invokeai/backend/quantization/scripts/quantize_t5_xxl_bnb_llm_int8.py

@@ -0,0 +1,92 @@
+from pathlib import Path
+
+import accelerate
+from safetensors.torch import load_file, save_file
+from transformers import AutoConfig, AutoModelForTextEncoding, T5EncoderModel
+
+from invokeai.backend.quantization.bnb_llm_int8 import quantize_model_llm_int8
+from invokeai.backend.quantization.scripts.load_flux_model_bnb_nf4 import log_time
+
+
+def load_state_dict_into_t5(model: T5EncoderModel, state_dict: dict):
+    # There is a shared reference to a single weight tensor in the model.
+    # Both "encoder.embed_tokens.weight" and "shared.weight" refer to the same tensor, so only the latter should
+    # be present in the state_dict.
+    missing_keys, unexpected_keys = model.load_state_dict(state_dict, strict=False, assign=True)
+    assert len(unexpected_keys) == 0
+    assert set(missing_keys) == {"encoder.embed_tokens.weight"}
+    # Assert that the layers we expect to be shared are actually shared.
+    assert model.encoder.embed_tokens.weight is model.shared.weight
+
+
+def main():
+    """A script for quantizing a T5 text encoder model using the bitsandbytes LLM.int8() quantization method.
+
+    This script is primarily intended for reference. The script params (e.g. the model_path, modules_to_not_convert,
+    etc.) are hardcoded and would need to be modified for other use cases.
+    """
+    model_path = Path("/data/misc/text_encoder_2")
+
+    with log_time("Intialize T5 on meta device"):
+        model_config = AutoConfig.from_pretrained(model_path)
+        with accelerate.init_empty_weights():
+            model = AutoModelForTextEncoding.from_config(model_config)
+
+    # TODO(ryand): We may want to add some modules to not quantize here (e.g. the proj_out layer). See the accelerate
+    # `get_keys_to_not_convert(...)` function for a heuristic to determine which modules to not quantize.
+    modules_to_not_convert: set[str] = set()
+
+    model_int8_path = model_path / "bnb_llm_int8.safetensors"
+    if model_int8_path.exists():
+        # The quantized model already exists, load it and return it.
+        print(f"A pre-quantized model already exists at '{model_int8_path}'. Attempting to load it...")
+
+        # Replace the linear layers with LLM.int8() quantized linear layers (still on the meta device).
+        with log_time("Replace linear layers with LLM.int8() layers"), accelerate.init_empty_weights():
+            model = quantize_model_llm_int8(model, modules_to_not_convert=modules_to_not_convert)
+
+        with log_time("Load state dict into model"):
+            sd = load_file(model_int8_path)
+            load_state_dict_into_t5(model, sd)
+
+        with log_time("Move model to cuda"):
+            model = model.to("cuda")
+
+        print(f"Successfully loaded pre-quantized model from '{model_int8_path}'.")
+
+    else:
+        # The quantized model does not exist, quantize the model and save it.
+        print(f"No pre-quantized model found at '{model_int8_path}'. Quantizing the model...")
+
+        with log_time("Replace linear layers with LLM.int8() layers"), accelerate.init_empty_weights():
+            model = quantize_model_llm_int8(model, modules_to_not_convert=modules_to_not_convert)
+
+        with log_time("Load state dict into model"):
+            # Load sharded state dict.
+            files = list(model_path.glob("*.safetensors"))
+            state_dict = {}
+            for file in files:
+                sd = load_file(file)
+                state_dict.update(sd)
+            load_state_dict_into_t5(model, state_dict)
+
+        with log_time("Move model to cuda and quantize"):
+            model = model.to("cuda")
+
+        with log_time("Save quantized model"):
+            model_int8_path.parent.mkdir(parents=True, exist_ok=True)
+            state_dict = model.state_dict()
+            state_dict.pop("encoder.embed_tokens.weight")
+            save_file(state_dict, model_int8_path)
+            # This handling of shared weights could also be achieved with save_model(...), but then we'd lose control
+            # over which keys are kept. And, the corresponding load_model(...) function does not support assign=True.
+            # save_model(model, model_int8_path)
+
+        print(f"Successfully quantized and saved model to '{model_int8_path}'.")
+
+    assert isinstance(model, T5EncoderModel)
+    return model
+
+
+if __name__ == "__main__":
+    main()

invokeai/backend/stable_diffusion/diffusion/conditioning_data.py

@@ -25,11 +25,6 @@ class BasicConditioningInfo:
         return self
 
 
-@dataclass
-class ConditioningFieldData:
-    conditionings: List[BasicConditioningInfo]
-
-
 @dataclass
 class SDXLConditioningInfo(BasicConditioningInfo):
     """SDXL text conditioning information produced by Compel."""
@@ -43,6 +38,22 @@ class SDXLConditioningInfo(BasicConditioningInfo):
         return super().to(device=device, dtype=dtype)
 
 
+@dataclass
+class FLUXConditioningInfo:
+    clip_embeds: torch.Tensor
+    t5_embeds: torch.Tensor
+
+    def to(self, device: torch.device | None = None, dtype: torch.dtype | None = None):
+        self.clip_embeds = self.clip_embeds.to(device=device, dtype=dtype)
+        self.t5_embeds = self.t5_embeds.to(device=device, dtype=dtype)
+        return self
+
+
+@dataclass
+class ConditioningFieldData:
+    conditionings: List[BasicConditioningInfo] | List[SDXLConditioningInfo] | List[FLUXConditioningInfo]
+
+
 @dataclass
 class IPAdapterConditioningInfo:
     cond_image_prompt_embeds: torch.Tensor

invokeai/backend/util/__init__.py

@@ -3,10 +3,9 @@ Initialization file for invokeai.backend.util
 """
 
 from invokeai.backend.util.logging import InvokeAILogger
-from invokeai.backend.util.util import GIG, Chdir, directory_size
+from invokeai.backend.util.util import Chdir, directory_size
 
 __all__ = [
-    "GIG",
     "directory_size",
     "Chdir",
     "InvokeAILogger",

invokeai/backend/util/util.py

@@ -7,9 +7,6 @@ from pathlib import Path
 
 from PIL import Image
 
-# actual size of a gig
-GIG = 1073741824
-
 
 def slugify(value: str, allow_unicode: bool = False) -> str:
     """

invokeai/frontend/web/.eslintrc.js

@@ -12,6 +12,10 @@ 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',
   },
   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,8 @@ 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',
   ],
   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.32",
     "@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-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"

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",
@@ -115,6 +116,7 @@
         "githubLabel": "Github",
         "goTo": "Go to",
         "hotkeysLabel": "Hotkeys",
+        "loadingImage": "Loading Image",
         "imageFailedToLoad": "Unable to Load Image",
         "img2img": "Image To Image",
         "inpaint": "inpaint",
@@ -162,10 +164,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 +327,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",
@@ -696,6 +706,8 @@
         "availableModels": "Available Models",
         "baseModel": "Base Model",
         "cancel": "Cancel",
+        "clipEmbed": "CLIP Embed",
+        "clipVision": "CLIP Vision",
         "config": "Config",
         "convert": "Convert",
         "convertingModelBegin": "Converting Model. Please wait.",
@@ -783,13 +795,16 @@
         "settings": "Settings",
         "simpleModelPlaceholder": "URL or path to a local file or diffusers folder",
         "source": "Source",
+        "spandrelImageToImage": "Image to Image (Spandrel)",
         "starterModels": "Starter Models",
+        "starterModelsInModelManager": "Starter Models can be found in Model Manager",
         "syncModels": "Sync Models",
         "textualInversions": "Textual Inversions",
         "triggerPhrases": "Trigger Phrases",
         "loraTriggerPhrases": "LoRA Trigger Phrases",
         "mainModelTriggerPhrases": "Main Model Trigger Phrases",
         "typePhraseHere": "Type phrase here",
+        "t5Encoder": "T5 Encoder",
         "upcastAttention": "Upcast Attention",
         "uploadImage": "Upload Image",
         "urlOrLocalPath": "URL or Local Path",
@@ -1095,7 +1110,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.",
@@ -1562,7 +1576,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",
@@ -1640,39 +1654,150 @@
         "storeNotInitialized": "Store is not initialized"
     },
     "controlLayers": {
-        "deleteAll": "Delete All",
+        "bookmark": "Bookmark for Quick Switch",
+        "removeBookmark": "Remove Bookmark",
+        "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",
         "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",
+        "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",
+            "preview": "Preview",
+            "apply": "Apply",
+            "cancel": "Cancel"
+        },
+        "transform": {
+            "transform": "Transform",
+            "fitToBbox": "Fit to Bbox",
+            "reset": "Reset",
+            "apply": "Apply",
+            "cancel": "Cancel"
+        }
     },
     "upscaling": {
         "upscale": "Upscale",
@@ -1760,5 +1885,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,12 +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';
@@ -39,11 +44,18 @@ interface Props {
     action: 'sendToImg2Img' | 'sendToCanvas' | 'useAllParameters';
   };
   selectedWorkflowId?: string;
-  destination?: InvokeTabName | undefined;
+  selectedStylePresetId?: string;
+  destination?: TabName;
 }
 
-const App = ({ config = DEFAULT_CONFIG, selectedImage, selectedWorkflowId, destination }: Props) => {
-  const language = useAppSelector(languageSelector);
+const App = ({
+  config = DEFAULT_CONFIG,
+  selectedImage,
+  selectedWorkflowId,
+  selectedStylePresetId,
+  destination,
+}: Props) => {
+  const language = useAppSelector(selectLanguage);
   const logger = useLogger('system');
   const dispatch = useAppDispatch();
   const clearStorage = useClearStorage();
@@ -81,6 +93,12 @@ const App = ({ config = DEFAULT_CONFIG, selectedImage, selectedWorkflowId, desti
     }
   }, [selectedWorkflowId, getAndLoadWorkflow]);
 
+  useEffect(() => {
+    if (selectedStylePresetId) {
+      dispatch(activeStylePresetIdChanged(selectedStylePresetId));
+    }
+  }, [dispatch, selectedStylePresetId]);
+
   useEffect(() => {
     if (destination) {
       dispatch(setActiveTab(destination));
@@ -93,6 +111,7 @@ const App = ({ config = DEFAULT_CONFIG, selectedImage, selectedWorkflowId, desti
 
   useStarterModelsToast();
   useSyncQueueStatus();
+  useScopeFocusWatcher();
 
   return (
     <ErrorBoundary onReset={handleReset} FallbackComponent={AppErrorBoundaryFallback}>
@@ -105,7 +124,7 @@ const App = ({ config = DEFAULT_CONFIG, selectedImage, selectedWorkflowId, desti
         {...dropzone.getRootProps()}
       >
         <input {...dropzone.getInputProps()} />
-        <InvokeTabs />
+        <AppContent />
         <AnimatePresence>
           {dropzone.isDragActive && isHandlingUpload && (
             <ImageUploadOverlay dropzone={dropzone} setIsHandlingUpload={setIsHandlingUpload} />
@@ -116,7 +135,10 @@ const App = ({ config = DEFAULT_CONFIG, selectedImage, selectedWorkflowId, desti
       <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';
@@ -45,7 +45,8 @@ interface Props extends PropsWithChildren {
     action: 'sendToImg2Img' | 'sendToCanvas' | 'useAllParameters';
   };
   selectedWorkflowId?: string;
-  destination?: InvokeTabName;
+  selectedStylePresetId?: string;
+  destination?: TabName;
   customStarUi?: CustomStarUi;
   socketOptions?: Partial<ManagerOptions & SocketOptions>;
   isDebugging?: boolean;
@@ -66,6 +67,7 @@ const InvokeAIUI = ({
   queueId,
   selectedImage,
   selectedWorkflowId,
+  selectedStylePresetId,
   destination,
   customStarUi,
   socketOptions,
@@ -227,6 +229,7 @@ const InvokeAIUI = ({
                 config={config}
                 selectedImage={selectedImage}
                 selectedWorkflowId={selectedWorkflowId}
+                selectedStylePresetId={selectedStylePresetId}
                 destination={destination}
               />
             </AppDndContext>

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 { 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,6 @@ 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 { 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 +26,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';
@@ -83,7 +63,6 @@ addGalleryImageClickedListener(startAppListening);
 addGalleryOffsetChangedListener(startAppListening);
 
 // User Invoked
-addEnqueueRequestedCanvasListener(startAppListening);
 addEnqueueRequestedNodes(startAppListening);
 addEnqueueRequestedLinear(startAppListening);
 addEnqueueRequestedUpscale(startAppListening);
@@ -91,31 +70,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 +118,4 @@ addAdHocPostProcessingRequestedListener(startAppListening);
 addDynamicPromptsListener(startAppListening);
 
 addSetDefaultSettingsListener(startAppListening);
-addControlAdapterPreprocessor(startAppListening);
+// addControlAdapterPreprocessor(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,37 @@
-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';
+  sessionStagingAreaImageAccepted,
+  sessionStagingAreaReset,
+} from 'features/controlLayers/store/canvasSessionSlice';
+import { rasterLayerAdded } from 'features/controlLayers/store/canvasSlice';
+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 { $lastCanvasProgressEvent } from 'services/events/setEventListeners';
+import { assert } from 'tsafe';
 
-const matcher = isAnyOf(commitStagingAreaImage, discardStagedImages, resetCanvas, setInitialCanvasImage);
+const log = logger('canvas');
 
-export const addCommitStagingAreaImageListener = (startAppListening: AppStartListening) => {
+export const addStagingListeners = (startAppListening: AppStartListening) => {
   startAppListening({
-    matcher,
-    effect: async (_, { dispatch, getState }) => {
-      const log = logger('canvas');
-      const state = getState();
-      const { batchIds } = state.canvas;
-
+    actionCreator: sessionStagingAreaReset,
+    effect: async (_, { dispatch }) => {
       try {
         const req = dispatch(
-          queueApi.endpoints.cancelByBatchIds.initiate({ batch_ids: batchIds }, { fixedCacheKey: 'cancelByBatchIds' })
+          queueApi.endpoints.cancelByBatchOrigin.initiate(
+            { origin: 'canvas' },
+            { fixedCacheKey: 'cancelByBatchOrigin' }
+          )
         );
         const { canceled } = await req.unwrap();
         req.reset();
+
+        $lastCanvasProgressEvent.set(null);
+
         if (canceled > 0) {
           log.debug(`Canceled ${canceled} canvas batches`);
           toast({
@@ -36,7 +40,6 @@ export const addCommitStagingAreaImageListener = (startAppListening: AppStartLis
             status: 'success',
           });
         }
-        dispatch(canvasBatchIdsReset());
       } catch {
         log.error('Failed to cancel canvas batches');
         toast({
@@ -47,4 +50,26 @@ export const addCommitStagingAreaImageListener = (startAppListening: AppStartLis
       }
     },
   });
+
+  startAppListening({
+    actionCreator: sessionStagingAreaImageAccepted,
+    effect: (action, api) => {
+      const { index } = action.payload;
+      const state = api.getState();
+      const stagingAreaImage = state.canvasSession.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(sessionStagingAreaReset());
+    },
+  });
 };

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/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,21 @@
+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/konva/CanvasManager';
+import { sessionStagingAreaReset, sessionStartedStaging } from 'features/controlLayers/store/canvasSessionSlice';
 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 +23,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 (!state.canvasSession.isStaging && state.canvasSession.sendToCanvas) {
+        dispatch(sessionStartedStaging());
+        didStartStaging = true;
       }
 
-      const batchConfig = prepareLinearUIBatch(state, graph, prepend);
+      const abortStaging = () => {
+        if (didStartStaging && getState().canvasSession.isStaging) {
+          dispatch(sessionStagingAreaReset());
+        }
+      };
+
+      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.canvasSession.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

@@ -14,9 +14,9 @@ export const addEnqueueRequestedUpscale = (startAppListening: AppStartListening)
       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, {

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 sessions?)
+
 // 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,19 @@
 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 {
-  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, isImageViewerOpenChanged, 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 +22,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 +38,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 +62,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 +78,38 @@ 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 imageObject = imageDTOToImageObject(activeData.payload.imageDTO);
+        const { x, y } = selectCanvasSlice(getState()).bbox.rect;
+        const overrides: Partial<CanvasControlLayerState> = {
+          objects: [imageObject],
+          position: { x, y },
+        };
+        dispatch(controlLayerAdded({ overrides, isSelected: 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,8 @@
 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 { 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 +10,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;
@@ -81,15 +71,6 @@ export const addImageUploadedFulfilledListener = (startAppListening: AppStartLis
         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 +80,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 +115,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/DocumentSize/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 })
+      );
+    });
   });
 };