src/sdxl_inversion_pipeline.py

# Code is based on ReNoise https://github.com/garibida/ReNoise-Inversion

import torch
from typing import Any, Callable, Dict, List, Optional, Tuple, Union

from diffusers import (
    StableDiffusionXLImg2ImgPipeline,
)
from diffusers.utils.torch_utils import randn_tensor

from diffusers.pipelines.stable_diffusion_xl.pipeline_stable_diffusion_xl import (
    StableDiffusionXLPipelineOutput,
    retrieve_timesteps,
    PipelineImageInput
)

from src.eunms import Epsilon_Update_Type


def _backward_ddim(x_tm1, alpha_t, alpha_tm1, eps_xt):
    """

    let a = alpha_t, b = alpha_{t - 1}

    We have a > b,

    x_{t} - x_{t - 1} = sqrt(a) ((sqrt(1/b) - sqrt(1/a)) * x_{t-1} + (sqrt(1/a - 1) - sqrt(1/b - 1)) * eps_{t-1})

    From https://arxiv.org/pdf/2105.05233.pdf, section F.

    """

    a, b = alpha_t, alpha_tm1
    sa = a ** 0.5
    sb = b ** 0.5

    return sa * ((1 / sb) * x_tm1 + ((1 / a - 1) ** 0.5 - (1 / b - 1) ** 0.5) * eps_xt)


class SDXLDDIMPipeline(StableDiffusionXLImg2ImgPipeline):
    # @torch.no_grad()
    def __call__(

            self,

            prompt: Union[str, List[str]] = None,

            prompt_2: Optional[Union[str, List[str]]] = None,

            image: PipelineImageInput = None,

            strength: float = 0.3,

            num_inversion_steps: int = 50,

            timesteps: List[int] = None,

            denoising_start: Optional[float] = None,

            denoising_end: Optional[float] = None,

            guidance_scale: float = 1.0,

            negative_prompt: Optional[Union[str, List[str]]] = None,

            negative_prompt_2: Optional[Union[str, List[str]]] = None,

            num_images_per_prompt: Optional[int] = 1,

            eta: float = 0.0,

            generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,

            latents: Optional[torch.FloatTensor] = None,

            prompt_embeds: Optional[torch.FloatTensor] = None,

            negative_prompt_embeds: Optional[torch.FloatTensor] = None,

            pooled_prompt_embeds: Optional[torch.FloatTensor] = None,

            negative_pooled_prompt_embeds: Optional[torch.FloatTensor] = None,

            ip_adapter_image: Optional[PipelineImageInput] = None,

            output_type: Optional[str] = "pil",

            return_dict: bool = True,

            cross_attention_kwargs: Optional[Dict[str, Any]] = None,

            guidance_rescale: float = 0.0,

            original_size: Tuple[int, int] = None,

            crops_coords_top_left: Tuple[int, int] = (0, 0),

            target_size: Tuple[int, int] = None,

            negative_original_size: Optional[Tuple[int, int]] = None,

            negative_crops_coords_top_left: Tuple[int, int] = (0, 0),

            negative_target_size: Optional[Tuple[int, int]] = None,

            aesthetic_score: float = 6.0,

            negative_aesthetic_score: float = 2.5,

            clip_skip: Optional[int] = None,

            callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None,

            callback_on_step_end_tensor_inputs: List[str] = ["latents"],

            num_inference_steps: int = 50,

            inv_hp=None,

            **kwargs,

    ):
        callback = kwargs.pop("callback", None)
        callback_steps = kwargs.pop("callback_steps", None)

        if callback is not None:
            deprecate(
                "callback",
                "1.0.0",
                "Passing `callback` as an input argument to `__call__` is deprecated, consider use `callback_on_step_end`",
            )
        if callback_steps is not None:
            deprecate(
                "callback_steps",
                "1.0.0",
                "Passing `callback_steps` as an input argument to `__call__` is deprecated, consider use `callback_on_step_end`",
            )

        # 1. Check inputs. Raise error if not correct
        self.check_inputs(
            prompt,
            prompt_2,
            strength,
            num_inversion_steps,
            callback_steps,
            negative_prompt,
            negative_prompt_2,
            prompt_embeds,
            negative_prompt_embeds,
            callback_on_step_end_tensor_inputs,
        )

        denoising_start_fr = 1.0 - denoising_start
        denoising_start = denoising_start

        self._guidance_scale = guidance_scale
        self._guidance_rescale = guidance_rescale
        self._clip_skip = clip_skip
        self._cross_attention_kwargs = cross_attention_kwargs
        self._denoising_end = denoising_end
        self._denoising_start = denoising_start

        # 2. Define call parameters
        if prompt is not None and isinstance(prompt, str):
            batch_size = 1
        elif prompt is not None and isinstance(prompt, list):
            batch_size = len(prompt)
        else:
            batch_size = prompt_embeds.shape[0]

        device = self._execution_device

        # 3. Encode input prompt
        text_encoder_lora_scale = (
            self.cross_attention_kwargs.get("scale", None) if self.cross_attention_kwargs is not None else None
        )
        (
            prompt_embeds,
            negative_prompt_embeds,
            pooled_prompt_embeds,
            negative_pooled_prompt_embeds,
        ) = self.encode_prompt(
            prompt=prompt,
            prompt_2=prompt_2,
            device=device,
            num_images_per_prompt=num_images_per_prompt,
            do_classifier_free_guidance=self.do_classifier_free_guidance,
            negative_prompt=negative_prompt,
            negative_prompt_2=negative_prompt_2,
            prompt_embeds=prompt_embeds,
            negative_prompt_embeds=negative_prompt_embeds,
            pooled_prompt_embeds=pooled_prompt_embeds,
            negative_pooled_prompt_embeds=negative_pooled_prompt_embeds,
            lora_scale=text_encoder_lora_scale,
            clip_skip=self.clip_skip,
        )

        # 4. Preprocess image
        image = self.image_processor.preprocess(image)

        # 5. Prepare timesteps
        def denoising_value_valid(dnv):
            return isinstance(self.denoising_end, float) and 0 < dnv < 1

        timesteps, num_inversion_steps = retrieve_timesteps(self.scheduler, num_inversion_steps, device, timesteps)
        timesteps_num_inference_steps, num_inference_steps = retrieve_timesteps(self.scheduler_inference,
                                                                                num_inference_steps, device, None)

        timesteps, num_inversion_steps = self.get_timesteps(
            num_inversion_steps,
            strength,
            device,
            denoising_start=self.denoising_start if denoising_value_valid else None,
        )
        # latent_timestep = timesteps[:1].repeat(batch_size * num_images_per_prompt)

        # add_noise = True if self.denoising_start is None else False
        # 6. Prepare latent variables
        with torch.no_grad():
            latents = self.prepare_latents(
                image,
                None,
                batch_size,
                num_images_per_prompt,
                prompt_embeds.dtype,
                device,
                generator,
                False,
            )
        # 7. Prepare extra step kwargs.
        extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta)

        height, width = latents.shape[-2:]
        height = height * self.vae_scale_factor
        width = width * self.vae_scale_factor

        original_size = original_size or (height, width)
        target_size = target_size or (height, width)

        # 8. Prepare added time ids & embeddings
        if negative_original_size is None:
            negative_original_size = original_size
        if negative_target_size is None:
            negative_target_size = target_size

        add_text_embeds = pooled_prompt_embeds
        if self.text_encoder_2 is None:
            text_encoder_projection_dim = int(pooled_prompt_embeds.shape[-1])
        else:
            text_encoder_projection_dim = self.text_encoder_2.config.projection_dim

        add_time_ids, add_neg_time_ids = self._get_add_time_ids(
            original_size,
            crops_coords_top_left,
            target_size,
            aesthetic_score,
            negative_aesthetic_score,
            negative_original_size,
            negative_crops_coords_top_left,
            negative_target_size,
            dtype=prompt_embeds.dtype,
            text_encoder_projection_dim=text_encoder_projection_dim,
        )
        add_time_ids = add_time_ids.repeat(batch_size * num_images_per_prompt, 1)

        if self.do_classifier_free_guidance:
            prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds], dim=0)
            add_text_embeds = torch.cat([negative_pooled_prompt_embeds, add_text_embeds], dim=0)
            add_neg_time_ids = add_neg_time_ids.repeat(batch_size * num_images_per_prompt, 1)
            add_time_ids = torch.cat([add_neg_time_ids, add_time_ids], dim=0)

        prompt_embeds = prompt_embeds.to(device)
        add_text_embeds = add_text_embeds.to(device)
        add_time_ids = add_time_ids.to(device)

        if ip_adapter_image is not None:
            image_embeds, negative_image_embeds = self.encode_image(ip_adapter_image, device, num_images_per_prompt)
            if self.do_classifier_free_guidance:
                image_embeds = torch.cat([negative_image_embeds, image_embeds])
                image_embeds = image_embeds.to(device)

        # 9. Denoising loop
        num_warmup_steps = max(len(timesteps) - num_inversion_steps * self.scheduler.order, 0)
        prev_timestep = None

        self._num_timesteps = len(timesteps)
        self.prev_z = torch.clone(latents)
        self.prev_z4 = torch.clone(latents)
        self.z_0 = torch.clone(latents)
        g_cpu = torch.Generator().manual_seed(7865)
        self.noise = randn_tensor(self.z_0.shape, generator=g_cpu, device=self.z_0.device, dtype=self.z_0.dtype)

        # Friendly inversion params
        timesteps_for = reversed(timesteps)
        noise = randn_tensor(latents.shape, generator=g_cpu, device=latents.device, dtype=latents.dtype)
        #latents = latents
        z_T = latents.clone()

        all_latents = [latents.clone()]
        with self.progress_bar(total=num_inversion_steps) as progress_bar:
            for i, t in enumerate(timesteps_for):

                added_cond_kwargs = {"text_embeds": add_text_embeds, "time_ids": add_time_ids}
                if ip_adapter_image is not None:
                    added_cond_kwargs["image_embeds"] = image_embeds

                z_tp1 = self.inversion_step(latents,
                                            t,
                                            prompt_embeds,
                                            added_cond_kwargs,
                                            prev_timestep=prev_timestep,
                                            inv_hp=inv_hp,
                                            z_0=self.z_0)

                prev_timestep = t
                latents = z_tp1

                all_latents.append(latents.clone())

                if callback_on_step_end is not None:
                    callback_kwargs = {}
                    for k in callback_on_step_end_tensor_inputs:
                        callback_kwargs[k] = locals()[k]
                    callback_outputs = callback_on_step_end(self, i, t, callback_kwargs)

                    latents = callback_outputs.pop("latents", latents)
                    prompt_embeds = callback_outputs.pop("prompt_embeds", prompt_embeds)
                    negative_prompt_embeds = callback_outputs.pop("negative_prompt_embeds", negative_prompt_embeds)
                    add_text_embeds = callback_outputs.pop("add_text_embeds", add_text_embeds)
                    negative_pooled_prompt_embeds = callback_outputs.pop(
                        "negative_pooled_prompt_embeds", negative_pooled_prompt_embeds
                    )
                    add_time_ids = callback_outputs.pop("add_time_ids", add_time_ids)
                    add_neg_time_ids = callback_outputs.pop("add_neg_time_ids", add_neg_time_ids)

                # call the callback, if provided
                if i == len(timesteps) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0):
                    progress_bar.update()
                    if callback is not None and i % callback_steps == 0:
                        step_idx = i // getattr(self.scheduler, "order", 1)
                        callback(step_idx, t, latents)

        image = latents

        # Offload all models
        self.maybe_free_model_hooks()

        return StableDiffusionXLPipelineOutput(images=image), all_latents

    def get_timestamp_dist(self, z_0, timesteps):
        timesteps = timesteps.to(z_0.device)
        if "cuda" in str(z_0.device):
            sigma = self.scheduler.sigmas.cuda()[:-1][self.scheduler.timesteps == timesteps]
        else:
            sigma = self.scheduler.sigmas[:-1][self.scheduler.timesteps == timesteps]
        z_0 = z_0.reshape(-1, 1)

        def gaussian_pdf(x):
            shape = x.shape
            x = x.reshape(-1, 1)
            all_probs = - 0.5 * torch.pow(((x - z_0) / sigma), 2)
            return all_probs.reshape(shape)

        return gaussian_pdf

    # @torch.no_grad()
    def inversion_step(

            self,

            z_t: torch.tensor,

            t: torch.tensor,

            prompt_embeds,

            added_cond_kwargs,

            prev_timestep: Optional[torch.tensor] = None,

            inv_hp=None,

            z_0=None,

    ) -> torch.tensor:

        n_iters, alpha, lr = inv_hp
        latent = z_t
        best_latent = None
        best_score = torch.inf
        curr_dist = self.get_timestamp_dist(z_0, t)
        for i in range(n_iters):
            latent.requires_grad = True
            noise_pred = self.unet_pass(latent, t, prompt_embeds, added_cond_kwargs)

            next_latent = self.backward_step(noise_pred, t, z_t, prev_timestep)
            f_x = (next_latent - latent).abs() - alpha * curr_dist(next_latent)
            score = f_x.mean()

            if score < best_score:
                best_score = score
                best_latent = next_latent.detach()

            f_x.sum().backward()
            latent = latent - lr * (f_x / latent.grad)
            latent.grad = None
            latent._grad_fn = None

        # if self.cfg.update_epsilon_type != Epsilon_Update_Type.NONE:
        #     noise_pred = self.unet_pass(best_latent, t, prompt_embeds, added_cond_kwargs)
        #     self.scheduler.step_and_update_noise(noise_pred, t, best_latent, z_t, return_dict=False,
        #                                          update_epsilon_type=self.cfg.update_epsilon_type)
        return best_latent

    @torch.no_grad()
    def unet_pass(self, z_t, t, prompt_embeds, added_cond_kwargs):
        latent_model_input = torch.cat([z_t] * 2) if self.do_classifier_free_guidance else z_t
        latent_model_input = self.scheduler.scale_model_input(latent_model_input, t)
        return self.unet(
            latent_model_input,
            t,
            encoder_hidden_states=prompt_embeds,
            timestep_cond=None,
            cross_attention_kwargs=self.cross_attention_kwargs,
            added_cond_kwargs=added_cond_kwargs,
            return_dict=False,
        )[0]

    @torch.no_grad()
    def backward_step(self, nosie_pred, t, z_t, prev_timestep):
        extra_step_kwargs = {}
        return self.scheduler.inv_step(nosie_pred, t, z_t, **extra_step_kwargs, return_dict=False)[0].detach()