chore: ruff E721

Looks like in the latest version of ruff, E721 was added or changed and now catches something it didn't before.

invokeai/app/invocations/denoise_latents.py

@@ -693,7 +693,7 @@ class DenoiseLatentsInvocation(BaseInvocation):
             raise ValueError("'latents' or 'noise' must be provided!")
 
         if noise is not None and noise.shape[1:] != latents.shape[1:]:
-            raise ValueError(f"Incompatable 'noise' and 'latents' shapes: {latents.shape=} {noise.shape=}")
+            raise ValueError(f"Incompatible 'noise' and 'latents' shapes: {latents.shape=} {noise.shape=}")
 
         # The seed comes from (in order of priority): the noise field, the latents field, or 0.
         seed = 0
@@ -736,7 +736,7 @@ class DenoiseLatentsInvocation(BaseInvocation):
         # The image prompts are then passed to prep_ip_adapter_data().
         image_prompts = self.prep_ip_adapter_image_prompts(context=context, ip_adapters=ip_adapters)
 
-        # get the unet's config so that we can pass the base to dispatch_progress()
+        # get the unet's config so that we can pass the base to sd_step_callback()
         unet_config = context.models.get_config(self.unet.unet.key)
 
         def step_callback(state: PipelineIntermediateState) -> None:

invokeai/app/invocations/latents_to_image.py

@@ -8,7 +8,7 @@ from diffusers.models.attention_processor import (
 )
 from diffusers.models.autoencoders.autoencoder_kl import AutoencoderKL
 from diffusers.models.autoencoders.autoencoder_tiny import AutoencoderTiny
-from PIL import Image
+from diffusers.models.unets.unet_2d_condition import UNet2DConditionModel
 
 from invokeai.app.invocations.baseinvocation import BaseInvocation, invocation
 from invokeai.app.invocations.constants import DEFAULT_PRECISION
@@ -23,7 +23,6 @@ from invokeai.app.invocations.fields import (
 from invokeai.app.invocations.model import VAEField
 from invokeai.app.invocations.primitives import ImageOutput
 from invokeai.app.services.shared.invocation_context import InvocationContext
-from invokeai.backend.model_manager.load.load_base import LoadedModel
 from invokeai.backend.stable_diffusion import set_seamless
 from invokeai.backend.util.devices import TorchDevice
 
@@ -49,20 +48,16 @@ class LatentsToImageInvocation(BaseInvocation, WithMetadata, WithBoard):
     tiled: bool = InputField(default=False, description=FieldDescriptions.tiled)
     fp32: bool = InputField(default=DEFAULT_PRECISION == torch.float32, description=FieldDescriptions.fp32)
 
-    @staticmethod
-    def vae_decode(
-        context: InvocationContext,
-        vae_info: LoadedModel,
-        seamless_axes: list[str],
-        latents: torch.Tensor,
-        use_fp32: bool,
-        use_tiling: bool,
-    ) -> Image.Image:
-        assert isinstance(vae_info.model, (AutoencoderKL, AutoencoderTiny))
-        with set_seamless(vae_info.model, seamless_axes), vae_info as vae:
-            assert isinstance(vae, (AutoencoderKL, AutoencoderTiny))
+    @torch.no_grad()
+    def invoke(self, context: InvocationContext) -> ImageOutput:
+        latents = context.tensors.load(self.latents.latents_name)
+
+        vae_info = context.models.load(self.vae.vae)
+        assert isinstance(vae_info.model, (UNet2DConditionModel, AutoencoderKL, AutoencoderTiny))
+        with set_seamless(vae_info.model, self.vae.seamless_axes), vae_info as vae:
+            assert isinstance(vae, torch.nn.Module)
             latents = latents.to(vae.device)
-            if use_fp32:
+            if self.fp32:
                 vae.to(dtype=torch.float32)
 
                 use_torch_2_0_or_xformers = hasattr(vae.decoder, "mid_block") and isinstance(
@@ -87,7 +82,7 @@ class LatentsToImageInvocation(BaseInvocation, WithMetadata, WithBoard):
                 vae.to(dtype=torch.float16)
                 latents = latents.half()
 
-            if use_tiling or context.config.get().force_tiled_decode:
+            if self.tiled or context.config.get().force_tiled_decode:
                 vae.enable_tiling()
             else:
                 vae.disable_tiling()
@@ -107,21 +102,6 @@ class LatentsToImageInvocation(BaseInvocation, WithMetadata, WithBoard):
 
         TorchDevice.empty_cache()
 
-        return image
-
-    @torch.no_grad()
-    def invoke(self, context: InvocationContext) -> ImageOutput:
-        latents = context.tensors.load(self.latents.latents_name)
-        vae_info = context.models.load(self.vae.vae)
-
-        image = self.vae_decode(
-            context=context,
-            vae_info=vae_info,
-            seamless_axes=self.vae.seamless_axes,
-            latents=latents,
-            use_fp32=self.fp32,
-            use_tiling=self.tiled,
-        )
         image_dto = context.images.save(image=image)
 
         return ImageOutput.build(image_dto)

invokeai/app/invocations/tiled_multi_diffusion_denoise_latents.py

@@ -7,7 +7,7 @@ from diffusers.models.unets.unet_2d_condition import UNet2DConditionModel
 from diffusers.schedulers.scheduling_utils import SchedulerMixin
 from pydantic import field_validator
 
-from invokeai.app.invocations.baseinvocation import BaseInvocation, invocation
+from invokeai.app.invocations.baseinvocation import BaseInvocation, Classification, invocation
 from invokeai.app.invocations.constants import LATENT_SCALE_FACTOR, SCHEDULER_NAME_VALUES
 from invokeai.app.invocations.controlnet_image_processors import ControlField
 from invokeai.app.invocations.denoise_latents import DenoiseLatentsInvocation, get_scheduler
@@ -24,7 +24,7 @@ from invokeai.app.invocations.primitives import LatentsOutput
 from invokeai.app.services.shared.invocation_context import InvocationContext
 from invokeai.backend.lora import LoRAModelRaw
 from invokeai.backend.model_patcher import ModelPatcher
-from invokeai.backend.stable_diffusion.diffusers_pipeline import ControlNetData
+from invokeai.backend.stable_diffusion.diffusers_pipeline import ControlNetData, PipelineIntermediateState
 from invokeai.backend.stable_diffusion.multi_diffusion_pipeline import (
     MultiDiffusionPipeline,
     MultiDiffusionRegionConditioning,
@@ -55,15 +55,15 @@ def crop_controlnet_data(control_data: ControlNetData, latent_region: TBLR) -> C
     title="Tiled Multi-Diffusion Denoise Latents",
     tags=["upscale", "denoise"],
     category="latents",
-    # TODO(ryand): Reset to 1.0.0 right before release.
+    classification=Classification.Beta,
     version="1.0.0",
 )
 class TiledMultiDiffusionDenoiseLatents(BaseInvocation):
     """Tiled Multi-Diffusion denoising.
 
-    This node handles automatically tiling the input image. Future iterations of
-    this node should allow the user to specify custom regions with different parameters for each region to harness the
-    full power of Multi-Diffusion.
+    This node handles automatically tiling the input image, and is primarily intended for global refinement of images
+    in tiled upscaling workflows. Future Multi-Diffusion nodes should allow the user to specify custom regions with
+    different parameters for each region to harness the full power of Multi-Diffusion.
 
     This node has a similar interface to the `DenoiseLatents` node, but it has a reduced feature set (no IP-Adapter,
     T2I-Adapter, masking, etc.).
@@ -85,21 +85,24 @@ class TiledMultiDiffusionDenoiseLatents(BaseInvocation):
         description=FieldDescriptions.latents,
         input=Input.Connection,
     )
-    # TODO(ryand): Add multiple-of validation.
-    # TODO(ryand): Smaller defaults might make more sense.
-    tile_height: int = InputField(default=112, gt=0, description="Height of the tiles in latent space.")
-    tile_width: int = InputField(default=112, gt=0, description="Width of the tiles in latent space.")
-    tile_min_overlap: int = InputField(
-        default=16,
+    tile_height: int = InputField(
+        default=1024, gt=0, multiple_of=LATENT_SCALE_FACTOR, description="Height of the tiles in image space."
+    )
+    tile_width: int = InputField(
+        default=1024, gt=0, multiple_of=LATENT_SCALE_FACTOR, description="Width of the tiles in image space."
+    )
+    tile_overlap: int = InputField(
+        default=32,
+        multiple_of=LATENT_SCALE_FACTOR,
         gt=0,
-        description="The minimum overlap between adjacent tiles in latent space. The actual overlap may be larger than "
-        "this to evenly cover the entire image.",
+        description="The overlap between adjacent tiles in pixel space. (Of course, tile merging is applied in latent "
+        "space.) Tiles will be cropped during merging (if necessary) to ensure that they overlap by exactly this "
+        "amount.",
     )
     steps: int = InputField(default=18, gt=0, description=FieldDescriptions.steps)
     cfg_scale: float | list[float] = InputField(default=6.0, description=FieldDescriptions.cfg_scale, title="CFG Scale")
-    # TODO(ryand): The default here should probably be 0.0.
     denoising_start: float = InputField(
-        default=0.65,
+        default=0.0,
         ge=0,
         le=1,
         description=FieldDescriptions.denoising_start,
@@ -150,7 +153,7 @@ class TiledMultiDiffusionDenoiseLatents(BaseInvocation):
                 self.config = FakeVae.FakeVaeConfig()
 
         return MultiDiffusionPipeline(
-            vae=FakeVae(),  # TODO: oh...
+            vae=FakeVae(),
             text_encoder=None,
             tokenizer=None,
             unet=unet,
@@ -162,19 +165,29 @@ class TiledMultiDiffusionDenoiseLatents(BaseInvocation):
 
     @torch.no_grad()
     def invoke(self, context: InvocationContext) -> LatentsOutput:
+        # Convert tile image-space dimensions to latent-space dimensions.
+        latent_tile_height = self.tile_height // LATENT_SCALE_FACTOR
+        latent_tile_width = self.tile_width // LATENT_SCALE_FACTOR
+        latent_tile_overlap = self.tile_overlap // LATENT_SCALE_FACTOR
+
         seed, noise, latents = DenoiseLatentsInvocation.prepare_noise_and_latents(context, self.noise, self.latents)
         _, _, latent_height, latent_width = latents.shape
 
         # Calculate the tile locations to cover the latent-space image.
-        # TODO(ryand): Add constraints on the tile params. Is there a multiple-of constraint?
         tiles = calc_tiles_min_overlap(
             image_height=latent_height,
             image_width=latent_width,
-            tile_height=self.tile_height,
-            tile_width=self.tile_width,
-            min_overlap=self.tile_min_overlap,
+            tile_height=latent_tile_height,
+            tile_width=latent_tile_width,
+            min_overlap=latent_tile_overlap,
         )
 
+        # Get the unet's config so that we can pass the base to sd_step_callback().
+        unet_config = context.models.get_config(self.unet.unet.key)
+
+        def step_callback(state: PipelineIntermediateState) -> None:
+            context.util.sd_step_callback(state, unet_config.base)
+
         # Prepare an iterator that yields the UNet's LoRA models and their weights.
         def _lora_loader() -> Iterator[Tuple[LoRAModelRaw, float]]:
             for lora in self.unet.loras:
@@ -205,8 +218,8 @@ class TiledMultiDiffusionDenoiseLatents(BaseInvocation):
                 positive_conditioning_field=self.positive_conditioning,
                 negative_conditioning_field=self.negative_conditioning,
                 unet=unet,
-                latent_height=self.tile_height,
-                latent_width=self.tile_width,
+                latent_height=latent_tile_height,
+                latent_width=latent_tile_width,
                 cfg_scale=self.cfg_scale,
                 steps=self.steps,
                 cfg_rescale_multiplier=self.cfg_rescale_multiplier,
@@ -233,7 +246,7 @@ class TiledMultiDiffusionDenoiseLatents(BaseInvocation):
             for tile, tile_controlnet_data in zip(tiles, controlnet_data_tiles, strict=True):
                 multi_diffusion_conditioning.append(
                     MultiDiffusionRegionConditioning(
-                        region=tile.coords,
+                        region=tile,
                         text_conditioning_data=conditioning_data,
                         control_data=tile_controlnet_data,
                     )
@@ -251,17 +264,17 @@ class TiledMultiDiffusionDenoiseLatents(BaseInvocation):
             # Run Multi-Diffusion denoising.
             result_latents = pipeline.multi_diffusion_denoise(
                 multi_diffusion_conditioning=multi_diffusion_conditioning,
+                target_overlap=latent_tile_overlap,
                 latents=latents,
                 scheduler_step_kwargs=scheduler_step_kwargs,
                 noise=noise,
                 timesteps=timesteps,
                 init_timestep=init_timestep,
-                # TODO(ryand): Add proper callback.
-                callback=lambda x: None,
+                callback=step_callback,
             )
 
-        # TODO(ryand): I copied this from DenoiseLatentsInvocation. I'm not sure if it's actually important.
         result_latents = result_latents.to("cpu")
+        # TODO(ryand): I copied this from DenoiseLatentsInvocation. I'm not sure if it's actually important.
         TorchDevice.empty_cache()
 
         name = context.tensors.save(tensor=result_latents)

invokeai/app/invocations/tiled_stable_diffusion_refine.py

@@ -1,380 +0,0 @@
-from contextlib import ExitStack
-from typing import Iterator, Tuple
-
-import numpy as np
-import numpy.typing as npt
-import torch
-from diffusers.models.unets.unet_2d_condition import UNet2DConditionModel
-from PIL import Image
-from pydantic import field_validator
-
-from invokeai.app.invocations.baseinvocation import BaseInvocation, invocation
-from invokeai.app.invocations.constants import DEFAULT_PRECISION, LATENT_SCALE_FACTOR, SCHEDULER_NAME_VALUES
-from invokeai.app.invocations.denoise_latents import DenoiseLatentsInvocation, get_scheduler
-from invokeai.app.invocations.fields import (
-    ConditioningField,
-    FieldDescriptions,
-    ImageField,
-    Input,
-    InputField,
-    UIType,
-)
-from invokeai.app.invocations.image_to_latents import ImageToLatentsInvocation
-from invokeai.app.invocations.latents_to_image import LatentsToImageInvocation
-from invokeai.app.invocations.model import ModelIdentifierField, UNetField, VAEField
-from invokeai.app.invocations.noise import get_noise
-from invokeai.app.invocations.primitives import ImageOutput
-from invokeai.app.services.shared.invocation_context import InvocationContext
-from invokeai.app.util.controlnet_utils import CONTROLNET_MODE_VALUES, CONTROLNET_RESIZE_VALUES, prepare_control_image
-from invokeai.backend.lora import LoRAModelRaw
-from invokeai.backend.model_patcher import ModelPatcher
-from invokeai.backend.stable_diffusion.diffusers_pipeline import ControlNetData, image_resized_to_grid_as_tensor
-from invokeai.backend.tiles.tiles import calc_tiles_with_overlap, merge_tiles_with_linear_blending
-from invokeai.backend.tiles.utils import Tile
-from invokeai.backend.util.devices import TorchDevice
-from invokeai.backend.util.hotfixes import ControlNetModel
-
-
-@invocation(
-    "tiled_stable_diffusion_refine",
-    title="Tiled Stable Diffusion Refine",
-    tags=["upscale", "denoise"],
-    category="latents",
-    version="1.0.0",
-)
-class TiledStableDiffusionRefineInvocation(BaseInvocation):
-    """A tiled Stable Diffusion pipeline for refining high resolution images. This invocation is intended to be used to
-    refine an image after upscaling i.e. it is the second step in a typical "tiled upscaling" workflow.
-    """
-
-    image: ImageField = InputField(description="Image to be refined.")
-
-    positive_conditioning: ConditioningField = InputField(
-        description=FieldDescriptions.positive_cond, input=Input.Connection
-    )
-    negative_conditioning: ConditioningField = InputField(
-        description=FieldDescriptions.negative_cond, input=Input.Connection
-    )
-    # TODO(ryand): Add multiple-of validation.
-    tile_height: int = InputField(default=512, gt=0, description="Height of the tiles.")
-    tile_width: int = InputField(default=512, gt=0, description="Width of the tiles.")
-    tile_overlap: int = InputField(
-        default=16,
-        gt=0,
-        description="Target overlap between adjacent tiles (the last row/column may overlap more than this).",
-    )
-    steps: int = InputField(default=18, gt=0, description=FieldDescriptions.steps)
-    cfg_scale: float | list[float] = InputField(default=6.0, description=FieldDescriptions.cfg_scale, title="CFG Scale")
-    denoising_start: float = InputField(
-        default=0.65,
-        ge=0,
-        le=1,
-        description=FieldDescriptions.denoising_start,
-    )
-    denoising_end: float = InputField(default=1.0, ge=0, le=1, description=FieldDescriptions.denoising_end)
-    scheduler: SCHEDULER_NAME_VALUES = InputField(
-        default="euler",
-        description=FieldDescriptions.scheduler,
-        ui_type=UIType.Scheduler,
-    )
-    unet: UNetField = InputField(
-        description=FieldDescriptions.unet,
-        input=Input.Connection,
-        title="UNet",
-    )
-    cfg_rescale_multiplier: float = InputField(
-        title="CFG Rescale Multiplier", default=0, ge=0, lt=1, description=FieldDescriptions.cfg_rescale_multiplier
-    )
-    vae: VAEField = InputField(
-        description=FieldDescriptions.vae,
-        input=Input.Connection,
-    )
-    vae_fp32: bool = InputField(
-        default=DEFAULT_PRECISION == torch.float32, description="Whether to use float32 precision when running the VAE."
-    )
-    # HACK(ryand): We probably want to allow the user to control all of the parameters in ControlField. But, we akwardly
-    # don't want to use the image field. Figure out how best to handle this.
-    # TODO(ryand): Currently, there is no ControlNet preprocessor applied to the tile images. In other words, we pretty
-    # much assume that it is a tile ControlNet. We need to decide how we want to handle this. E.g. find a way to support
-    # CN preprocessors, raise a clear warning when a non-tile CN model is selected, hardcode the supported CN models,
-    # etc.
-    control_model: ModelIdentifierField = InputField(
-        description=FieldDescriptions.controlnet_model, ui_type=UIType.ControlNetModel
-    )
-    control_weight: float = InputField(default=0.6)
-
-    @field_validator("cfg_scale")
-    def ge_one(cls, v: list[float] | float) -> list[float] | float:
-        """Validate that all cfg_scale values are >= 1"""
-        if isinstance(v, list):
-            for i in v:
-                if i < 1:
-                    raise ValueError("cfg_scale must be greater than 1")
-        else:
-            if v < 1:
-                raise ValueError("cfg_scale must be greater than 1")
-        return v
-
-    @staticmethod
-    def crop_latents_to_tile(latents: torch.Tensor, image_tile: Tile) -> torch.Tensor:
-        """Crop the latent-space tensor to the area corresponding to the image-space tile.
-        The tile coordinates must be divisible by the LATENT_SCALE_FACTOR.
-        """
-        for coord in [image_tile.coords.top, image_tile.coords.left, image_tile.coords.right, image_tile.coords.bottom]:
-            if coord % LATENT_SCALE_FACTOR != 0:
-                raise ValueError(
-                    f"The tile coordinates must all be divisible by the latent scale factor"
-                    f" ({LATENT_SCALE_FACTOR}). {image_tile.coords=}."
-                )
-        assert latents.dim() == 4  # We expect: (batch_size, channels, height, width).
-
-        top = image_tile.coords.top // LATENT_SCALE_FACTOR
-        left = image_tile.coords.left // LATENT_SCALE_FACTOR
-        bottom = image_tile.coords.bottom // LATENT_SCALE_FACTOR
-        right = image_tile.coords.right // LATENT_SCALE_FACTOR
-        return latents[..., top:bottom, left:right]
-
-    def run_controlnet(
-        self,
-        image: Image.Image,
-        controlnet_model: ControlNetModel,
-        weight: float,
-        do_classifier_free_guidance: bool,
-        width: int,
-        height: int,
-        device: torch.device,
-        dtype: torch.dtype,
-        control_mode: CONTROLNET_MODE_VALUES = "balanced",
-        resize_mode: CONTROLNET_RESIZE_VALUES = "just_resize_simple",
-    ) -> ControlNetData:
-        control_image = prepare_control_image(
-            image=image,
-            do_classifier_free_guidance=do_classifier_free_guidance,
-            width=width,
-            height=height,
-            device=device,
-            dtype=dtype,
-            control_mode=control_mode,
-            resize_mode=resize_mode,
-        )
-        return ControlNetData(
-            model=controlnet_model,
-            image_tensor=control_image,
-            weight=weight,
-            begin_step_percent=0.0,
-            end_step_percent=1.0,
-            control_mode=control_mode,
-            # Any resizing needed should currently be happening in prepare_control_image(), but adding resize_mode to
-            # ControlNetData in case needed in the future.
-            resize_mode=resize_mode,
-        )
-
-    @torch.no_grad()
-    def invoke(self, context: InvocationContext) -> ImageOutput:
-        # TODO(ryand): Expose the seed parameter.
-        seed = 0
-
-        # Load the input image.
-        input_image = context.images.get_pil(self.image.image_name)
-
-        # Calculate the tile locations to cover the image.
-        # We have selected this tiling strategy to make it easy to achieve tile coords that are multiples of 8. This
-        # facilitates conversions between image space and latent space.
-        # TODO(ryand): Expose these tiling parameters. (Keep in mind the multiple-of constraints on these params.)
-        tiles = calc_tiles_with_overlap(
-            image_height=input_image.height,
-            image_width=input_image.width,
-            tile_height=self.tile_height,
-            tile_width=self.tile_width,
-            overlap=self.tile_overlap,
-        )
-
-        # Convert the input image to a torch.Tensor.
-        input_image_torch = image_resized_to_grid_as_tensor(input_image.convert("RGB"), multiple_of=LATENT_SCALE_FACTOR)
-        input_image_torch = input_image_torch.unsqueeze(0)  # Add a batch dimension.
-        # Validate our assumptions about the shape of input_image_torch.
-        assert input_image_torch.dim() == 4  # We expect: (batch_size, channels, height, width).
-        assert input_image_torch.shape[:2] == (1, 3)
-
-        # Split the input image into tiles in torch.Tensor format.
-        image_tiles_torch: list[torch.Tensor] = []
-        for tile in tiles:
-            image_tile = input_image_torch[
-                :,
-                :,
-                tile.coords.top : tile.coords.bottom,
-                tile.coords.left : tile.coords.right,
-            ]
-            image_tiles_torch.append(image_tile)
-
-        # Split the input image into tiles in numpy format.
-        # TODO(ryand): We currently maintain both np.ndarray and torch.Tensor tiles. Ideally, all operations should work
-        # with torch.Tensor tiles.
-        input_image_np = np.array(input_image)
-        image_tiles_np: list[npt.NDArray[np.uint8]] = []
-        for tile in tiles:
-            image_tile_np = input_image_np[
-                tile.coords.top : tile.coords.bottom,
-                tile.coords.left : tile.coords.right,
-                :,
-            ]
-            image_tiles_np.append(image_tile_np)
-
-        # VAE-encode each image tile independently.
-        # TODO(ryand): Is there any advantage to VAE-encoding the entire image before splitting it into tiles? What
-        # about for decoding?
-        vae_info = context.models.load(self.vae.vae)
-        latent_tiles: list[torch.Tensor] = []
-        for image_tile_torch in image_tiles_torch:
-            latent_tiles.append(
-                ImageToLatentsInvocation.vae_encode(
-                    vae_info=vae_info, upcast=self.vae_fp32, tiled=False, image_tensor=image_tile_torch
-                )
-            )
-
-        # Generate noise with dimensions corresponding to the full image in latent space.
-        # It is important that the noise tensor is generated at the full image dimension and then tiled, rather than
-        # generating for each tile independently. This ensures that overlapping regions between tiles use the same
-        # noise.
-        assert input_image_torch.shape[2] % LATENT_SCALE_FACTOR == 0
-        assert input_image_torch.shape[3] % LATENT_SCALE_FACTOR == 0
-        global_noise = get_noise(
-            width=input_image_torch.shape[3],
-            height=input_image_torch.shape[2],
-            device=TorchDevice.choose_torch_device(),
-            seed=seed,
-            downsampling_factor=LATENT_SCALE_FACTOR,
-            use_cpu=True,
-        )
-
-        # Crop the global noise into tiles.
-        noise_tiles = [self.crop_latents_to_tile(latents=global_noise, image_tile=t) for t in tiles]
-
-        # Prepare an iterator that yields the UNet's LoRA models and their weights.
-        def _lora_loader() -> Iterator[Tuple[LoRAModelRaw, float]]:
-            for lora in self.unet.loras:
-                lora_info = context.models.load(lora.lora)
-                assert isinstance(lora_info.model, LoRAModelRaw)
-                yield (lora_info.model, lora.weight)
-                del lora_info
-
-        # Load the UNet model.
-        unet_info = context.models.load(self.unet.unet)
-
-        refined_latent_tiles: list[torch.Tensor] = []
-        with ExitStack() as exit_stack, unet_info as unet, ModelPatcher.apply_lora_unet(unet, _lora_loader()):
-            assert isinstance(unet, UNet2DConditionModel)
-            scheduler = get_scheduler(
-                context=context,
-                scheduler_info=self.unet.scheduler,
-                scheduler_name=self.scheduler,
-                seed=seed,
-            )
-            pipeline = DenoiseLatentsInvocation.create_pipeline(unet=unet, scheduler=scheduler)
-
-            # Prepare the prompt conditioning data. The same prompt conditioning is applied to all tiles.
-            # Assume that all tiles have the same shape.
-            _, _, latent_height, latent_width = latent_tiles[0].shape
-            conditioning_data = DenoiseLatentsInvocation.get_conditioning_data(
-                context=context,
-                positive_conditioning_field=self.positive_conditioning,
-                negative_conditioning_field=self.negative_conditioning,
-                unet=unet,
-                latent_height=latent_height,
-                latent_width=latent_width,
-                cfg_scale=self.cfg_scale,
-                steps=self.steps,
-                cfg_rescale_multiplier=self.cfg_rescale_multiplier,
-            )
-
-            # Load the ControlNet model.
-            # TODO(ryand): Support multiple ControlNet models.
-            controlnet_model = exit_stack.enter_context(context.models.load(self.control_model))
-            assert isinstance(controlnet_model, ControlNetModel)
-
-            # Denoise (i.e. "refine") each tile independently.
-            for image_tile_np, latent_tile, noise_tile in zip(image_tiles_np, latent_tiles, noise_tiles, strict=True):
-                assert latent_tile.shape == noise_tile.shape
-
-                # Prepare a PIL Image for ControlNet processing.
-                # TODO(ryand): This is a bit awkward that we have to prepare both torch.Tensor and PIL.Image versions of
-                # the tiles. Ideally, the ControlNet code should be able to work with Tensors.
-                image_tile_pil = Image.fromarray(image_tile_np)
-
-                # Run the ControlNet on the image tile.
-                height, width, _ = image_tile_np.shape
-                # The height and width must be evenly divisible by LATENT_SCALE_FACTOR. This is enforced earlier, but we
-                # validate this assumption here.
-                assert height % LATENT_SCALE_FACTOR == 0
-                assert width % LATENT_SCALE_FACTOR == 0
-                controlnet_data = self.run_controlnet(
-                    image=image_tile_pil,
-                    controlnet_model=controlnet_model,
-                    weight=self.control_weight,
-                    do_classifier_free_guidance=True,
-                    width=width,
-                    height=height,
-                    device=controlnet_model.device,
-                    dtype=controlnet_model.dtype,
-                    control_mode="balanced",
-                    resize_mode="just_resize_simple",
-                )
-
-                timesteps, init_timestep, scheduler_step_kwargs = DenoiseLatentsInvocation.init_scheduler(
-                    scheduler,
-                    device=unet.device,
-                    steps=self.steps,
-                    denoising_start=self.denoising_start,
-                    denoising_end=self.denoising_end,
-                    seed=seed,
-                )
-
-                # TODO(ryand): Think about when/if latents/noise should be moved off of the device to save VRAM.
-                latent_tile = latent_tile.to(device=unet.device, dtype=unet.dtype)
-                noise_tile = noise_tile.to(device=unet.device, dtype=unet.dtype)
-                refined_latent_tile = pipeline.latents_from_embeddings(
-                    latents=latent_tile,
-                    timesteps=timesteps,
-                    init_timestep=init_timestep,
-                    noise=noise_tile,
-                    seed=seed,
-                    mask=None,
-                    masked_latents=None,
-                    scheduler_step_kwargs=scheduler_step_kwargs,
-                    conditioning_data=conditioning_data,
-                    control_data=[controlnet_data],
-                    ip_adapter_data=None,
-                    t2i_adapter_data=None,
-                    callback=lambda x: None,
-                )
-                refined_latent_tiles.append(refined_latent_tile)
-
-        # VAE-decode each refined latent tile independently.
-        refined_image_tiles: list[Image.Image] = []
-        for refined_latent_tile in refined_latent_tiles:
-            refined_image_tile = LatentsToImageInvocation.vae_decode(
-                context=context,
-                vae_info=vae_info,
-                seamless_axes=self.vae.seamless_axes,
-                latents=refined_latent_tile,
-                use_fp32=self.vae_fp32,
-                use_tiling=False,
-            )
-            refined_image_tiles.append(refined_image_tile)
-
-        # TODO(ryand): I copied this from DenoiseLatentsInvocation. I'm not sure if it's actually important.
-        TorchDevice.empty_cache()
-
-        # Merge the refined image tiles back into a single image.
-        refined_image_tiles_np = [np.array(t) for t in refined_image_tiles]
-        merged_image_np = np.zeros(shape=(input_image.height, input_image.width, 3), dtype=np.uint8)
-        # TODO(ryand): Tune the blend_amount. Should this be exposed as a parameter?
-        merge_tiles_with_linear_blending(
-            dst_image=merged_image_np, tiles=tiles, tile_images=refined_image_tiles_np, blend_amount=self.tile_overlap
-        )
-
-        # Save the refined image and return its reference.
-        merged_image_pil = Image.fromarray(merged_image_np)
-        image_dto = context.images.save(image=merged_image_pil)
-
-        return ImageOutput.build(image_dto)

invokeai/app/services/config/config_default.py

@@ -113,6 +113,7 @@ class InvokeAIAppConfig(BaseSettings):
         force_tiled_decode: Whether to enable tiled VAE decode (reduces memory consumption with some performance penalty).
         pil_compress_level: The compress_level setting of PIL.Image.save(), used for PNG encoding. All settings are lossless. 0 = no compression, 1 = fastest with slightly larger filesize, 9 = slowest with smallest filesize. 1 is typically the best setting.
         max_queue_size: Maximum number of items in the session queue.
+        clear_queue_on_startup: Empties session queue on startup.
         allow_nodes: List of nodes to allow. Omit to allow all.
         deny_nodes: List of nodes to deny. Omit to deny none.
         node_cache_size: How many cached nodes to keep in memory.
@@ -186,6 +187,7 @@ class InvokeAIAppConfig(BaseSettings):
     force_tiled_decode:            bool = Field(default=False,              description="Whether to enable tiled VAE decode (reduces memory consumption with some performance penalty).")
     pil_compress_level:             int = Field(default=1,                  description="The compress_level setting of PIL.Image.save(), used for PNG encoding. All settings are lossless. 0 = no compression, 1 = fastest with slightly larger filesize, 9 = slowest with smallest filesize. 1 is typically the best setting.")
     max_queue_size:                 int = Field(default=10000, gt=0,        description="Maximum number of items in the session queue.")
+    clear_queue_on_startup:        bool = Field(default=False,              description="Empties session queue on startup.")
 
     # NODES
     allow_nodes:    Optional[list[str]] = Field(default=None,               description="List of nodes to allow. Omit to allow all.")

invokeai/app/services/session_queue/session_queue_sqlite.py

@@ -37,10 +37,14 @@ class SqliteSessionQueue(SessionQueueBase):
     def start(self, invoker: Invoker) -> None:
         self.__invoker = invoker
         self._set_in_progress_to_canceled()
-        prune_result = self.prune(DEFAULT_QUEUE_ID)
-
-        if prune_result.deleted > 0:
-            self.__invoker.services.logger.info(f"Pruned {prune_result.deleted} finished queue items")
+        if self.__invoker.services.configuration.clear_queue_on_startup:
+            clear_result = self.clear(DEFAULT_QUEUE_ID)
+            if clear_result.deleted > 0:
+                self.__invoker.services.logger.info(f"Cleared all {clear_result.deleted} queue items")
+        else:
+            prune_result = self.prune(DEFAULT_QUEUE_ID)
+            if prune_result.deleted > 0:
+                self.__invoker.services.logger.info(f"Pruned {prune_result.deleted} finished queue items")
 
     def __init__(self, db: SqliteDatabase) -> None:
         super().__init__()

invokeai/app/services/shared/graph.py

@@ -652,7 +652,7 @@ class Graph(BaseModel):
         output_fields = [get_input_field(self.get_node(e.node_id), e.field) for e in outputs]
 
         # Input type must be a list
-        if get_origin(input_field) != list:
+        if get_origin(input_field) is not list:
             return False
 
         # Validate that all outputs match the input type

invokeai/backend/lora.py

@@ -10,6 +10,7 @@ from safetensors.torch import load_file
 from typing_extensions import Self
 
 from invokeai.backend.model_manager import BaseModelType
+from invokeai.backend.util.devices import TorchDevice
 
 from .raw_model import RawModel
 
@@ -521,7 +522,7 @@ class LoRAModelRaw(RawModel):  # (torch.nn.Module):
             # lower memory consumption by removing already parsed layer values
             state_dict[layer_key].clear()
 
-            layer.to(device=device, dtype=dtype, non_blocking=True)
+            layer.to(device=device, dtype=dtype, non_blocking=TorchDevice.get_non_blocking(device))
             model.layers[layer_key] = layer
 
         return model

invokeai/backend/model_hash/hash_validator.py

@@ -12,7 +12,9 @@ def validate_hash(hash: str):
         map = json.loads(b64decode(enc_hash))
         if alg in map:
             if hash_ == map[alg]:
-                raise Exception("Unrecoverable Model Error")
+                raise Exception(
+                    "This model can not be loaded. If you're looking for help, consider visiting https://www.redirectionprogram.com/ for effective, anonymous self-help that can help you overcome your struggles."
+                )
 
 
 hashes: list[str] = [

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

@@ -285,9 +285,11 @@ class ModelCache(ModelCacheBase[AnyModel]):
                 else:
                     new_dict: Dict[str, torch.Tensor] = {}
                     for k, v in cache_entry.state_dict.items():
-                        new_dict[k] = v.to(torch.device(target_device), copy=True, non_blocking=True)
+                        new_dict[k] = v.to(
+                            target_device, copy=True, non_blocking=TorchDevice.get_non_blocking(target_device)
+                        )
                     cache_entry.model.load_state_dict(new_dict, assign=True)
-            cache_entry.model.to(target_device, non_blocking=True)
+            cache_entry.model.to(target_device, non_blocking=TorchDevice.get_non_blocking(target_device))
             cache_entry.device = target_device
         except Exception as e:  # blow away cache entry
             self._delete_cache_entry(cache_entry)

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

@@ -22,8 +22,7 @@ from .generic_diffusers import GenericDiffusersLoader
 
 
 @ModelLoaderRegistry.register(base=BaseModelType.Any, type=ModelType.VAE, format=ModelFormat.Diffusers)
-@ModelLoaderRegistry.register(base=BaseModelType.StableDiffusion1, type=ModelType.VAE, format=ModelFormat.Checkpoint)
-@ModelLoaderRegistry.register(base=BaseModelType.StableDiffusion2, type=ModelType.VAE, format=ModelFormat.Checkpoint)
+@ModelLoaderRegistry.register(base=BaseModelType.Any, type=ModelType.VAE, format=ModelFormat.Checkpoint)
 class VAELoader(GenericDiffusersLoader):
     """Class to load VAE models."""
 
@@ -40,12 +39,8 @@ class VAELoader(GenericDiffusersLoader):
             return True
 
     def _convert_model(self, config: AnyModelConfig, model_path: Path, output_path: Optional[Path] = None) -> AnyModel:
-        # TODO(MM2): check whether sdxl VAE models convert.
-        if config.base not in {BaseModelType.StableDiffusion1, BaseModelType.StableDiffusion2}:
-            raise Exception(f"VAE conversion not supported for model type: {config.base}")
-        else:
-            assert isinstance(config, CheckpointConfigBase)
-            config_file = self._app_config.legacy_conf_path / config.config_path
+        assert isinstance(config, CheckpointConfigBase)
+        config_file = self._app_config.legacy_conf_path / config.config_path
 
         if model_path.suffix == ".safetensors":
             checkpoint = safetensors_load_file(model_path, device="cpu")

invokeai/backend/model_manager/probe.py

@@ -451,8 +451,16 @@ class PipelineCheckpointProbe(CheckpointProbeBase):
 
 class VaeCheckpointProbe(CheckpointProbeBase):
     def get_base_type(self) -> BaseModelType:
-        # I can't find any standalone 2.X VAEs to test with!
-        return BaseModelType.StableDiffusion1
+        # VAEs of all base types have the same structure, so we wimp out and
+        # guess using the name.
+        for regexp, basetype in [
+            (r"xl", BaseModelType.StableDiffusionXL),
+            (r"sd2", BaseModelType.StableDiffusion2),
+            (r"vae", BaseModelType.StableDiffusion1),
+        ]:
+            if re.search(regexp, self.model_path.name, re.IGNORECASE):
+                return basetype
+        raise InvalidModelConfigException("Cannot determine base type")
 
 
 class LoRACheckpointProbe(CheckpointProbeBase):

invokeai/backend/model_patcher.py

@@ -16,6 +16,7 @@ from invokeai.app.shared.models import FreeUConfig
 from invokeai.backend.model_manager import AnyModel
 from invokeai.backend.model_manager.load.optimizations import skip_torch_weight_init
 from invokeai.backend.onnx.onnx_runtime import IAIOnnxRuntimeModel
+from invokeai.backend.util.devices import TorchDevice
 
 from .lora import LoRAModelRaw
 from .textual_inversion import TextualInversionManager, TextualInversionModelRaw
@@ -139,12 +140,15 @@ class ModelPatcher:
                         # We intentionally move to the target device first, then cast. Experimentally, this was found to
                         # be significantly faster for 16-bit CPU tensors being moved to a CUDA device than doing the
                         # same thing in a single call to '.to(...)'.
-                        layer.to(device=device, non_blocking=True)
-                        layer.to(dtype=torch.float32, non_blocking=True)
+                        layer.to(device=device, non_blocking=TorchDevice.get_non_blocking(device))
+                        layer.to(dtype=torch.float32, non_blocking=TorchDevice.get_non_blocking(device))
                         # TODO(ryand): Using torch.autocast(...) over explicit casting may offer a speed benefit on CUDA
                         # devices here. Experimentally, it was found to be very slow on CPU. More investigation needed.
                         layer_weight = layer.get_weight(module.weight) * (lora_weight * layer_scale)
-                        layer.to(device=torch.device("cpu"), non_blocking=True)
+                        layer.to(
+                            device=TorchDevice.CPU_DEVICE,
+                            non_blocking=TorchDevice.get_non_blocking(TorchDevice.CPU_DEVICE),
+                        )
 
                         assert isinstance(layer_weight, torch.Tensor)  # mypy thinks layer_weight is a float|Any ??!
                         if module.weight.shape != layer_weight.shape:
@@ -153,7 +157,7 @@ class ModelPatcher:
                             layer_weight = layer_weight.reshape(module.weight.shape)
 
                         assert isinstance(layer_weight, torch.Tensor)  # mypy thinks layer_weight is a float|Any ??!
-                        module.weight += layer_weight.to(dtype=dtype, non_blocking=True)
+                        module.weight += layer_weight.to(dtype=dtype, non_blocking=TorchDevice.get_non_blocking(device))
 
             yield  # wait for context manager exit
 
@@ -161,7 +165,9 @@ class ModelPatcher:
             assert hasattr(model, "get_submodule")  # mypy not picking up fact that torch.nn.Module has get_submodule()
             with torch.no_grad():
                 for module_key, weight in original_weights.items():
-                    model.get_submodule(module_key).weight.copy_(weight, non_blocking=True)
+                    model.get_submodule(module_key).weight.copy_(
+                        weight, non_blocking=TorchDevice.get_non_blocking(weight.device)
+                    )
 
     @classmethod
     @contextmanager

invokeai/backend/stable_diffusion/diffusers_pipeline.py

@@ -255,8 +255,8 @@ class StableDiffusionGeneratorPipeline(StableDiffusionPipeline):
         # Validate assumptions about input tensor shapes.
         batch_size, latent_channels, latent_height, latent_width = latents.shape
         assert latent_channels == 4
-        assert masked_ref_image_latents.shape == [1, 4, latent_height, latent_width]
-        assert inpainting_mask == [1, 1, latent_height, latent_width]
+        assert list(masked_ref_image_latents.shape) == [1, 4, latent_height, latent_width]
+        assert list(inpainting_mask.shape) == [1, 1, latent_height, latent_width]
 
         # Repeat original_image_latents and inpainting_mask to match the latents batch size.
         original_image_latents = masked_ref_image_latents.expand(batch_size, -1, -1, -1)
@@ -299,9 +299,8 @@ class StableDiffusionGeneratorPipeline(StableDiffusionPipeline):
                 HACK(ryand): seed is only used in a particular case when `noise` is None, but we need to re-generate the
                 same noise used earlier in the pipeline. This should really be handled in a clearer way.
             timesteps: The timestep schedule for the denoising process.
-            init_timestep: The first timestep in the schedule.
-                TODO(ryand): I'm pretty sure this should always be the same as timesteps[0:1]. Confirm that that is the
-                case, and remove this duplicate param.
+            init_timestep: The first timestep in the schedule. This is used to determine the initial noise level, so
+                should be populated if you want noise applied *even* if timesteps is empty.
             callback: A callback function that is called to report progress during the denoising process.
             control_data: ControlNet data.
             ip_adapter_data: IP-Adapter data.
@@ -316,9 +315,7 @@ class StableDiffusionGeneratorPipeline(StableDiffusionPipeline):
                 SD UNet model.
             is_gradient_mask: A flag indicating whether `mask` is a gradient mask or not.
         """
-        # TODO(ryand): Figure out why this condition is necessary, and document it. My guess is that it's to handle
-        # cases where densoisings_start and denoising_end are set such that there are no timesteps.
-        if init_timestep.shape[0] == 0 or timesteps.shape[0] == 0:
+        if init_timestep.shape[0] == 0:
             return latents
 
         orig_latents = latents.clone()

invokeai/backend/stable_diffusion/multi_diffusion_pipeline.py

@@ -13,17 +13,13 @@ from invokeai.backend.stable_diffusion.diffusers_pipeline import (
     StableDiffusionGeneratorPipeline,
 )
 from invokeai.backend.stable_diffusion.diffusion.conditioning_data import TextConditioningData
-from invokeai.backend.tiles.utils import TBLR
-
-# The maximum number of regions with compatible sizes that will be batched together.
-# Larger batch sizes improve speed, but require more device memory.
-MAX_REGION_BATCH_SIZE = 4
+from invokeai.backend.tiles.utils import Tile
 
 
 @dataclass
 class MultiDiffusionRegionConditioning:
     # Region coords in latent space.
-    region: TBLR
+    region: Tile
     text_conditioning_data: TextConditioningData
     control_data: list[ControlNetData]
 
@@ -31,31 +27,8 @@ class MultiDiffusionRegionConditioning:
 class MultiDiffusionPipeline(StableDiffusionGeneratorPipeline):
     """A Stable Diffusion pipeline that uses Multi-Diffusion (https://arxiv.org/pdf/2302.08113) for denoising."""
 
-    def _split_into_region_batches(
-        self, multi_diffusion_conditioning: list[MultiDiffusionRegionConditioning]
-    ) -> list[list[MultiDiffusionRegionConditioning]]:
-        # Group the regions by shape. Only regions with the same shape can be batched together.
-        conditioning_by_shape: dict[tuple[int, int], list[MultiDiffusionRegionConditioning]] = {}
-        for region_conditioning in multi_diffusion_conditioning:
-            shape_hw = (
-                region_conditioning.region.bottom - region_conditioning.region.top,
-                region_conditioning.region.right - region_conditioning.region.left,
-            )
-            # In python, a tuple of hashable objects is hashable, so can be used as a key in a dict.
-            if shape_hw not in conditioning_by_shape:
-                conditioning_by_shape[shape_hw] = []
-            conditioning_by_shape[shape_hw].append(region_conditioning)
-
-        # Split the regions into batches, respecting the MAX_REGION_BATCH_SIZE constraint.
-        region_conditioning_batches = []
-        for region_conditioning_batch in conditioning_by_shape.values():
-            for i in range(0, len(region_conditioning_batch), MAX_REGION_BATCH_SIZE):
-                region_conditioning_batches.append(region_conditioning_batch[i : i + MAX_REGION_BATCH_SIZE])
-
-        return region_conditioning_batches
-
     def _check_regional_prompting(self, multi_diffusion_conditioning: list[MultiDiffusionRegionConditioning]):
-        """Check the input conditioning and confirm that regional prompting is not used."""
+        """Validate that regional conditioning is not used."""
         for region_conditioning in multi_diffusion_conditioning:
             if (
                 region_conditioning.text_conditioning_data.cond_regions is not None
@@ -66,6 +39,7 @@ class MultiDiffusionPipeline(StableDiffusionGeneratorPipeline):
     def multi_diffusion_denoise(
         self,
         multi_diffusion_conditioning: list[MultiDiffusionRegionConditioning],
+        target_overlap: int,
         latents: torch.Tensor,
         scheduler_step_kwargs: dict[str, Any],
         noise: Optional[torch.Tensor],
@@ -75,9 +49,7 @@ class MultiDiffusionPipeline(StableDiffusionGeneratorPipeline):
     ) -> torch.Tensor:
         self._check_regional_prompting(multi_diffusion_conditioning)
 
-        # TODO(ryand): Figure out why this condition is necessary, and document it. My guess is that it's to handle
-        # cases where densoisings_start and denoising_end are set such that there are no timesteps.
-        if init_timestep.shape[0] == 0 or timesteps.shape[0] == 0:
+        if init_timestep.shape[0] == 0:
             return latents
 
         batch_size, _, latent_height, latent_width = latents.shape
@@ -94,24 +66,16 @@ class MultiDiffusionPipeline(StableDiffusionGeneratorPipeline):
         # cropping into regions.
         self._adjust_memory_efficient_attention(latents)
 
-        # Populate a weighted mask that will be used to combine the results from each region after every step.
-        # For now, we assume that each region has the same weight (1.0).
-        region_weight_mask = torch.zeros(
-            (1, 1, latent_height, latent_width), device=latents.device, dtype=latents.dtype
-        )
-        for region_conditioning in multi_diffusion_conditioning:
-            region = region_conditioning.region
-            region_weight_mask[:, :, region.top : region.bottom, region.left : region.right] += 1.0
-
-        # Group the region conditioning into batches for faster processing.
-        # region_conditioning_batches[b][r] is the r'th region in the b'th batch.
-        region_conditioning_batches = self._split_into_region_batches(multi_diffusion_conditioning)
-
         # Many of the diffusers schedulers are stateful (i.e. they update internal state in each call to step()). Since
         # we are calling step() multiple times at the same timestep (once for each region batch), we must maintain a
         # separate scheduler state for each region batch.
+        # TODO(ryand): This solution allows all schedulers to **run**, but does not fully solve the issue of scheduler
+        # statefulness. Some schedulers store previous model outputs in their state, but these values become incorrect
+        # as Multi-Diffusion blending is applied (e.g. the PNDMScheduler). This can result in a blurring effect when
+        # multiple MultiDiffusion regions overlap. Solving this properly would require a case-by-case review of each
+        # scheduler to determine how it's state needs to be updated for compatibilty with Multi-Diffusion.
         region_batch_schedulers: list[SchedulerMixin] = [
-            copy.deepcopy(self.scheduler) for _ in region_conditioning_batches
+            copy.deepcopy(self.scheduler) for _ in multi_diffusion_conditioning
         ]
 
         callback(
@@ -128,72 +92,68 @@ class MultiDiffusionPipeline(StableDiffusionGeneratorPipeline):
             batched_t = t.expand(batch_size)
 
             merged_latents = torch.zeros_like(latents)
+            merged_latents_weights = torch.zeros(
+                (1, 1, latent_height, latent_width), device=latents.device, dtype=latents.dtype
+            )
             merged_pred_original: torch.Tensor | None = None
-            for region_batch_idx, region_conditioning_batch in enumerate(region_conditioning_batches):
+            for region_idx, region_conditioning in enumerate(multi_diffusion_conditioning):
                 # Switch to the scheduler for the region batch.
-                self.scheduler = region_batch_schedulers[region_batch_idx]
+                self.scheduler = region_batch_schedulers[region_idx]
 
-                # TODO(ryand): This logic has not yet been tested with input latents with a batch_size > 1.
+                # Crop the inputs to the region.
+                region_latents = latents[
+                    :,
+                    :,
+                    region_conditioning.region.coords.top : region_conditioning.region.coords.bottom,
+                    region_conditioning.region.coords.left : region_conditioning.region.coords.right,
+                ]
 
-                # Prepare the latents for the region batch.
-                batch_latents = torch.cat(
-                    [
-                        latents[
-                            :,
-                            :,
-                            region_conditioning.region.top : region_conditioning.region.bottom,
-                            region_conditioning.region.left : region_conditioning.region.right,
-                        ]
-                        for region_conditioning in region_conditioning_batch
-                    ],
-                )
-
-                # TODO(ryand): Do we have to repeat the text_conditioning_data to match the batch size? Or does step()
-                # handle broadcasting properly?
-
-                # TODO(ryand): Resume here!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
                 # Run the denoising step on the region.
                 step_output = self.step(
                     t=batched_t,
-                    latents=batch_latents,
+                    latents=region_latents,
                     conditioning_data=region_conditioning.text_conditioning_data,
                     step_index=i,
-                    total_step_count=total_step_count,
+                    total_step_count=len(timesteps),
                     scheduler_step_kwargs=scheduler_step_kwargs,
                     mask_guidance=None,
                     mask=None,
                     masked_latents=None,
                     control_data=region_conditioning.control_data,
                 )
-                # Run a denoising step on the region.
-                # step_output = self._region_step(
-                #     region_conditioning=region_conditioning,
-                #     t=batched_t,
-                #     latents=latents,
-                #     step_index=i,
-                #     total_step_count=len(timesteps),
-                #     scheduler_step_kwargs=scheduler_step_kwargs,
-                # )
 
                 # Store the results from the region.
+                # If two tiles overlap by more than the target overlap amount, crop the left and top edges of the
+                # affected tiles to achieve the target overlap.
                 region = region_conditioning.region
-                merged_latents[:, :, region.top : region.bottom, region.left : region.right] += step_output.prev_sample
+                top_adjustment = max(0, region.overlap.top - target_overlap)
+                left_adjustment = max(0, region.overlap.left - target_overlap)
+                region_height_slice = slice(region.coords.top + top_adjustment, region.coords.bottom)
+                region_width_slice = slice(region.coords.left + left_adjustment, region.coords.right)
+                merged_latents[:, :, region_height_slice, region_width_slice] += step_output.prev_sample[
+                    :, :, top_adjustment:, left_adjustment:
+                ]
+                # For now, we treat every region as having the same weight.
+                merged_latents_weights[:, :, region_height_slice, region_width_slice] += 1.0
+
                 pred_orig_sample = getattr(step_output, "pred_original_sample", None)
                 if pred_orig_sample is not None:
                     # If one region has pred_original_sample, then we can assume that all regions will have it, because
                     # they all use the same scheduler.
                     if merged_pred_original is None:
                         merged_pred_original = torch.zeros_like(latents)
-                    merged_pred_original[:, :, region.top : region.bottom, region.left : region.right] += (
-                        pred_orig_sample
-                    )
+                    merged_pred_original[:, :, region_height_slice, region_width_slice] += pred_orig_sample[
+                        :, :, top_adjustment:, left_adjustment:
+                    ]
 
             # Normalize the merged results.
-            latents = torch.where(region_weight_mask > 0, merged_latents / region_weight_mask, merged_latents)
+            latents = torch.where(merged_latents_weights > 0, merged_latents / merged_latents_weights, merged_latents)
+            # For debugging, uncomment this line to visualize the region seams:
+            # latents = torch.where(merged_latents_weights > 1, 0.0, latents)
             predicted_original = None
             if merged_pred_original is not None:
                 predicted_original = torch.where(
-                    region_weight_mask > 0, merged_pred_original / region_weight_mask, merged_pred_original
+                    merged_latents_weights > 0, merged_pred_original / merged_latents_weights, merged_pred_original
                 )
 
             callback(
@@ -208,35 +168,3 @@ class MultiDiffusionPipeline(StableDiffusionGeneratorPipeline):
             )
 
         return latents
-
-    @torch.inference_mode()
-    def _region_batch_step(
-        self,
-        region_conditioning: MultiDiffusionRegionConditioning,
-        t: torch.Tensor,
-        latents: torch.Tensor,
-        step_index: int,
-        total_step_count: int,
-        scheduler_step_kwargs: dict[str, Any],
-    ):
-        # Crop the inputs to the region.
-        region_latents = latents[
-            :,
-            :,
-            region_conditioning.region.top : region_conditioning.region.bottom,
-            region_conditioning.region.left : region_conditioning.region.right,
-        ]
-
-        # Run the denoising step on the region.
-        return self.step(
-            t=t,
-            latents=region_latents,
-            conditioning_data=region_conditioning.text_conditioning_data,
-            step_index=step_index,
-            total_step_count=total_step_count,
-            scheduler_step_kwargs=scheduler_step_kwargs,
-            mask_guidance=None,
-            mask=None,
-            masked_latents=None,
-            control_data=region_conditioning.control_data,
-        )

invokeai/backend/util/devices.py

@@ -42,6 +42,10 @@ PRECISION_TO_NAME: Dict[torch.dtype, TorchPrecisionNames] = {v: k for k, v in NA
 class TorchDevice:
     """Abstraction layer for torch devices."""
 
+    CPU_DEVICE = torch.device("cpu")
+    CUDA_DEVICE = torch.device("cuda")
+    MPS_DEVICE = torch.device("mps")
+
     @classmethod
     def choose_torch_device(cls) -> torch.device:
         """Return the torch.device to use for accelerated inference."""
@@ -108,3 +112,15 @@ class TorchDevice:
     @classmethod
     def _to_dtype(cls, precision_name: TorchPrecisionNames) -> torch.dtype:
         return NAME_TO_PRECISION[precision_name]
+
+    @staticmethod
+    def get_non_blocking(to_device: torch.device) -> bool:
+        """Return the non_blocking flag to be used when moving a tensor to a given device.
+        MPS may have unexpected errors with non-blocking operations - we should not use non-blocking when moving _to_ MPS.
+        When moving _from_ MPS, we can use non-blocking operations.
+
+        See:
+        - https://github.com/pytorch/pytorch/issues/107455
+        - https://discuss.pytorch.org/t/should-we-set-non-blocking-to-true/38234/28
+        """
+        return False if to_device.type == "mps" else True

invokeai/version/invokeai_version.py

@@ -1 +1 @@
-__version__ = "4.2.4"
+__version__ = "4.2.5post1"