InvokeAI/invokeai/app/invocations/tiled_stable_diffusion_refine.py

import numpy as np
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 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.backend.stable_diffusion.diffusers_pipeline import image_resized_to_grid_as_tensor
from invokeai.backend.tiles.tiles import calc_tiles_min_overlap, merge_tiles_with_linear_blending
from invokeai.backend.tiles.utils import Tile
from invokeai.backend.util.devices import TorchDevice


@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
    )
    steps: int = InputField(default=10, gt=0, description=FieldDescriptions.steps)
    cfg_scale: float | list[float] = InputField(default=7.5, description=FieldDescriptions.cfg_scale, title="CFG Scale")
    denoising_start: float = InputField(
        default=0.0,
        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",
    )
    # control: Optional[Union[ControlField, list[ControlField]]] = InputField(
    #     default=None,
    #     input=Input.Connection,
    # )
    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."
    )

    @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]

    @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)
        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.

        # Calculate the tile locations to cover the image.
        # TODO(ryand): Expose these tiling parameters. (Keep in mind the multiple-of constraints on these params.)
        tiles = calc_tiles_min_overlap(
            image_height=input_image.height,
            image_width=input_image.width,
            tile_height=512,
            tile_width=512,
            min_overlap=128,
        )

        # 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.
        image_tiles: 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.append(image_tile)

        # 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 in image_tiles:
            latent_tiles.append(
                ImageToLatentsInvocation.vae_encode(
                    vae_info=vae_info, upcast=self.vae_fp32, tiled=False, image_tensor=image_tile
                )
            )

        # 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]

        # Load the UNet model.
        unet_info = context.models.load(self.unet.unet)

        refined_latent_tiles: list[torch.Tensor] = []
        with unet_info as unet:
            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,
            )

            # Denoise (i.e. "refine") each tile independently.
            for latent_tile, noise_tile in zip(latent_tiles, noise_tiles, strict=True):
                assert latent_tile.shape == noise_tile.shape

                num_inference_steps, 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,
                    gradient_mask=None,
                    num_inference_steps=num_inference_steps,
                    scheduler_step_kwargs=scheduler_step_kwargs,
                    conditioning_data=conditioning_data,
                    control_data=None,
                    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): Expose the blend_amount parameter, or set it based on the value of min_overlap used earlier.
        merge_tiles_with_linear_blending(
            dst_image=merged_image_np, tiles=tiles, tile_images=refined_image_tiles_np, blend_amount=32
        )

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