InvokeAI/invokeai/backend/stable_diffusion/multi_diffusion_pipeline.py

from __future__ import annotations

import copy
from dataclasses import dataclass
from typing import Any, Callable, Optional

import torch
from diffusers.schedulers.scheduling_utils import SchedulerMixin

from invokeai.backend.stable_diffusion.diffusers_pipeline import (
    ControlNetData,
    PipelineIntermediateState,
    StableDiffusionGeneratorPipeline,
)
from invokeai.backend.stable_diffusion.diffusion.conditioning_data import TextConditioningData
from invokeai.backend.tiles.utils import TBLR


@dataclass
class MultiDiffusionRegionConditioning:
    # Region coords in latent space.
    region: TBLR
    text_conditioning_data: TextConditioningData
    control_data: list[ControlNetData]


class MultiDiffusionPipeline(StableDiffusionGeneratorPipeline):
    """A Stable Diffusion pipeline that uses Multi-Diffusion (https://arxiv.org/pdf/2302.08113) for denoising."""

    def multi_diffusion_denoise(
        self,
        multi_diffusion_conditioning: list[MultiDiffusionRegionConditioning],
        latents: torch.Tensor,
        scheduler_step_kwargs: dict[str, Any],
        noise: Optional[torch.Tensor],
        timesteps: torch.Tensor,
        init_timestep: torch.Tensor,
        callback: Callable[[PipelineIntermediateState], None],
    ) -> torch.Tensor:
        # 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:
            return latents

        batch_size, _, latent_height, latent_width = latents.shape
        batched_init_timestep = init_timestep.expand(batch_size)

        # noise can be None if the latents have already been noised (e.g. when running the SDXL refiner).
        if noise is not None:
            # TODO(ryand): I'm pretty sure we should be applying init_noise_sigma in cases where we are starting with
            # full noise. Investigate the history of why this got commented out.
            # latents = noise * self.scheduler.init_noise_sigma # it's like in t2l according to diffusers
            latents = self.scheduler.add_noise(latents, noise, batched_init_timestep)

        # TODO(ryand): Look into the implications of passing in latents here that are larger than they will be after
        # 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 regions 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

        # 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.
        region_batch_schedulers: list[SchedulerMixin] = [
            copy.deepcopy(self.scheduler) for _ in multi_diffusion_conditioning
        ]

        callback(
            PipelineIntermediateState(
                step=-1,
                order=self.scheduler.order,
                total_steps=len(timesteps),
                timestep=self.scheduler.config.num_train_timesteps,
                latents=latents,
            )
        )

        for i, t in enumerate(self.progress_bar(timesteps)):
            batched_t = t.expand(batch_size)

            merged_latents = torch.zeros_like(latents)
            merged_pred_original: torch.Tensor | None = None
            for region_idx, region_conditioning in enumerate(multi_diffusion_conditioning):
                # Switch to the scheduler for the region batch.
                self.scheduler = region_batch_schedulers[region_idx]

                # 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.
                region = region_conditioning.region
                merged_latents[:, :, region.top : region.bottom, region.left : region.right] += step_output.prev_sample
                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
                    )

            # Normalize the merged results.
            latents = torch.where(region_weight_mask > 0, merged_latents / region_weight_mask, merged_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
                )

            callback(
                PipelineIntermediateState(
                    step=i,
                    order=self.scheduler.order,
                    total_steps=len(timesteps),
                    timestep=int(t),
                    latents=latents,
                    predicted_original=predicted_original,
                )
            )

        return latents

    @torch.inference_mode()
    def _region_step(
        self,
        region_conditioning: MultiDiffusionRegionConditioning,
        t: torch.Tensor,
        latents: torch.Tensor,
        step_index: int,
        total_step_count: int,
        scheduler_step_kwargs: dict[str, Any],
    ):
        use_regional_prompting = (
            region_conditioning.text_conditioning_data.cond_regions is not None
            or region_conditioning.text_conditioning_data.uncond_regions is not None
        )
        if use_regional_prompting:
            raise NotImplementedError("Regional prompting is not yet supported in Multi-Diffusion.")

        # 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,
        )