analogy_input_processor/pipeline.py

"""
ADOBE CONFIDENTIAL
Copyright 2024 Adobe
All Rights Reserved.
NOTICE: All information contained herein is, and remains
the property of Adobe and its suppliers, if any. The intellectual
and technical concepts contained herein are proprietary to Adobe
and its suppliers and are protected by all applicable intellectual
property laws, including trade secret and copyright laws.
Dissemination of this information or reproduction of this material
is strictly forbidden unless prior written permission is obtained
from Adobe.
"""

from typing import Callable, List, Optional, Union
import inspect
import einops
import PIL.Image
import numpy as np
import torch as th
import torch.nn as nn
from torchvision import transforms

from diffusers import ModelMixin
from transformers import AutoModel, AutoConfig, SiglipVisionConfig, Dinov2Config, Dinov2Model
from transformers import SiglipVisionModel
from diffusers import DiffusionPipeline
from diffusers.image_processor import VaeImageProcessor
from diffusers.models import AutoencoderKL, UNet2DConditionModel
from diffusers.schedulers import KarrasDiffusionSchedulers
from diffusers.utils.torch_utils import randn_tensor
from diffusers.pipelines.pipeline_utils import DiffusionPipeline, ImagePipelineOutput

from diffusers.configuration_utils import ConfigMixin, register_to_config
# REf: https://github.com/tatp22/multidim-positional-encoding/tree/master

    
OUT_SIZE = 768
IN_SIZE = 2048

DINO_SIZE = 224
DINO_MEAN = [0.485, 0.456, 0.406]
DINO_STD = [0.229, 0.224, 0.225]

SIGLIP_SIZE = 256
SIGLIP_MEAN = [0.5]
SIGLIP_STD = [0.5]


def get_emb(sin_inp):
    """
    Gets a base embedding for one dimension with sin and cos intertwined
    """
    emb = th.stack((sin_inp.sin(), sin_inp.cos()), dim=-1)
    return th.flatten(emb, -2, -1)


class PositionalEncoding1D(nn.Module):
    def __init__(self, channels):
        """
        :param channels: The last dimension of the tensor you want to apply pos emb to.
        """
        super(PositionalEncoding1D, self).__init__()
        self.org_channels = channels
        channels = int(np.ceil(channels / 2) * 2)
        self.channels = channels
        inv_freq = 1.0 / (10000 ** (th.arange(0, channels, 2).float() / channels))
        self.register_buffer("inv_freq", inv_freq)
        self.register_buffer("cached_penc", None, persistent=False)

    def forward(self, tensor):
        """
        :param tensor: A 3d tensor of size (batch_size, x, ch)
        :return: Positional Encoding Matrix of size (batch_size, x, ch)
        """
        if len(tensor.shape) != 3:
            raise RuntimeError("The input tensor has to be 3d!")

        if self.cached_penc is not None and self.cached_penc.shape == tensor.shape:
            return self.cached_penc

        self.cached_penc = None
        batch_size, x, orig_ch = tensor.shape
        pos_x = th.arange(x, device=tensor.device, dtype=self.inv_freq.dtype)
        sin_inp_x = th.einsum("i,j->ij", pos_x, self.inv_freq)
        emb_x = get_emb(sin_inp_x)
        emb = th.zeros((x, self.channels), device=tensor.device, dtype=tensor.dtype)
        emb[:, : self.channels] = emb_x

        self.cached_penc = emb[None, :, :orig_ch].repeat(batch_size, 1, 1)
        return self.cached_penc



class PositionalEncoding3D(nn.Module):
    def __init__(self, channels):
        """
        :param channels: The last dimension of the tensor you want to apply pos emb to.
        """
        super(PositionalEncoding3D, self).__init__()
        self.org_channels = channels
        channels = int(np.ceil(channels / 6) * 2)
        if channels % 2:
            channels += 1
        self.channels = channels
        inv_freq = 1.0 / (10000 ** (th.arange(0, channels, 2).float() / channels))
        self.register_buffer("inv_freq", inv_freq)
        self.register_buffer("cached_penc", None, persistent=False)

    def forward(self, tensor):
        """
        :param tensor: A 5d tensor of size (batch_size, x, y, z, ch)
        :return: Positional Encoding Matrix of size (batch_size, x, y, z, ch)
        """
        if len(tensor.shape) != 5:
            raise RuntimeError("The input tensor has to be 5d!")

        if self.cached_penc is not None and self.cached_penc.shape == tensor.shape:
            return self.cached_penc

        self.cached_penc = None
        batch_size, x, y, z, orig_ch = tensor.shape
        pos_x = th.arange(x, device=tensor.device, dtype=self.inv_freq.dtype)
        pos_y = th.arange(y, device=tensor.device, dtype=self.inv_freq.dtype)
        pos_z = th.arange(z, device=tensor.device, dtype=self.inv_freq.dtype)
        sin_inp_x = th.einsum("i,j->ij", pos_x, self.inv_freq)
        sin_inp_y = th.einsum("i,j->ij", pos_y, self.inv_freq)
        sin_inp_z = th.einsum("i,j->ij", pos_z, self.inv_freq)
        emb_x = get_emb(sin_inp_x).unsqueeze(1).unsqueeze(1)
        emb_y = get_emb(sin_inp_y).unsqueeze(1)
        emb_z = get_emb(sin_inp_z)
        emb = th.zeros(
            (x, y, z, self.channels * 3),
            device=tensor.device,
            dtype=tensor.dtype,
        )
        emb[:, :, :, : self.channels] = emb_x
        emb[:, :, :, self.channels : 2 * self.channels] = emb_y
        emb[:, :, :, 2 * self.channels :] = emb_z

        self.cached_penc = emb[None, :, :, :, :orig_ch].repeat(batch_size, 1, 1, 1, 1)
        return self.cached_penc

class AnalogyInputProcessor(ModelMixin, ConfigMixin):
    
    @register_to_config
    def __init__(self,):
        super(AnalogyInputProcessor, self).__init__()
        
        self.dino_transform = transforms.Compose(
            [
                transforms.Resize((DINO_SIZE, DINO_SIZE)),
                transforms.ToTensor(),
                transforms.Normalize(DINO_MEAN, DINO_STD), # SIGLIP normalization
            ]
        )
        
        self.siglip_transform = transforms.Compose(
            [
                transforms.Resize((SIGLIP_SIZE, SIGLIP_SIZE)),
                transforms.ToTensor(),
                transforms.Normalize(SIGLIP_MEAN, SIGLIP_STD), # SIGLIP normalization
            ]
        )
        
        dino_mean = th.tensor(DINO_MEAN).view(1, 3, 1, 1)
        dino_std = th.tensor(DINO_STD).view(1, 3, 1, 1)
        siglip_mean = [SIGLIP_MEAN[0],] * 3
        siglip_std = [SIGLIP_STD[0],] * 3
        siglip_mean = th.tensor(siglip_mean).view(1, 3, 1, 1)
        siglip_std = th.tensor(siglip_std).view(1, 3, 1, 1)
        self.register_buffer("dino_mean", dino_mean)
        self.register_buffer("dino_std", dino_std)
        self.register_buffer("siglip_mean", siglip_mean)
        self.register_buffer("siglip_std", siglip_std)
        
    def __call__(self, analogy_prompt):
        # List of tuples of (A, A*, B)
        img_a_dino = []
        img_a_siglip = []
        img_a_star_dino = []
        img_a_star_siglip = []
        img_b_dino = []
        img_b_siglip = []
        
        for im_set in analogy_prompt:
            img_a, img_a_star, img_b = im_set
            img_a_dino.append(self.dino_transform(img_a))
            img_a_siglip.append(self.siglip_transform(img_a))
            img_a_star_dino.append(self.dino_transform(img_a_star))
            img_a_star_siglip.append(self.siglip_transform(img_a_star))
            img_b_dino.append(self.dino_transform(img_b))
            img_b_siglip.append(self.siglip_transform(img_b))
        
        img_a_dino = th.stack(img_a_dino, 0)
        img_a_siglip = th.stack(img_a_siglip, 0)
        img_a_star_dino = th.stack(img_a_star_dino, 0)
        img_a_star_siglip = th.stack(img_a_star_siglip, 0)
        img_b_dino = th.stack(img_b_dino, 0)
        img_b_siglip = th.stack(img_b_siglip, 0)
        
        dino_combined_input = th.stack([img_b_dino, img_a_dino, img_a_star_dino], 0)
        siglip_combined_input = th.stack([img_b_siglip, img_a_siglip, img_a_star_siglip], 0)
        
        return dino_combined_input, siglip_combined_input
    def get_negative(self, dino_in, siglip_in):
        
        dino_i = ((dino_in * 0 + 0.5) - self.dino_mean) / self.dino_std
        siglip_i = ((siglip_in * 0 + 0.5) - self.siglip_mean) / self.siglip_std
        return dino_i, siglip_i
            

class AnalogyProjector(ModelMixin, ConfigMixin):
    
    @register_to_config
    def __init__(self):
        super(AnalogyProjector, self).__init__()
        self.projector = DinoSiglipMixer()
        self.pos_embd_1D = PositionalEncoding1D(OUT_SIZE)
        self.pos_embd_3D = PositionalEncoding3D(OUT_SIZE)
        

    def forward(self, dino_in, siglip_in, batch_size):
        
        image_embeddings = self.projector(dino_in, siglip_in)
        
        image_embeddings = einops.rearrange(image_embeddings, '(k b) t d -> b k t d', b=batch_size)
        image_embeddings = self.position_embd(image_embeddings)
        return image_embeddings

    def position_embd(self, image_embeddings, concat=False):
        canvas_embd = image_embeddings[:, :, 1:, :]
        batch_size = canvas_embd.shape[0]
        type_size = canvas_embd.shape[1]
        xy_size = canvas_embd.shape[2]
        
        x_size = int(xy_size ** 0.5)

        canvas_embd = canvas_embd.reshape(batch_size, type_size, x_size, x_size, -1)
        if concat:
            canvas_embd = th.cat([canvas_embd, self.pos_embd_3D(canvas_embd)], -1)
        else:
            canvas_embd = self.pos_embd_3D(canvas_embd) + canvas_embd
        canvas_embd = canvas_embd.reshape(batch_size, type_size, xy_size, -1)

        class_embd = image_embeddings[:, :, 0, :]
        if concat:
            class_embd = th.cat([class_embd, self.pos_embd_1D(class_embd)], -1)
        else:    
            class_embd = self.pos_embd_1D(class_embd) + class_embd
        all_embd_list = []
        for i in range(type_size):
            all_embd_list.append(class_embd[:, i:i+1])
            all_embd_list.append(canvas_embd[:, i])
        image_embeddings = th.cat(all_embd_list, 1)
        return image_embeddings


class HighLowMixer(th.nn.Module):
    def __init__(self, in_size=IN_SIZE, out_size=OUT_SIZE):
        super().__init__()
        mid_size = (in_size + out_size) // 2
        
        self.lower_projector = th.nn.Sequential(
            th.nn.LayerNorm(IN_SIZE//2),
            th.nn.SiLU()
        )
        self.upper_projector = th.nn.Sequential(
            th.nn.LayerNorm(IN_SIZE//2),
            th.nn.SiLU()
        )
        self.projectors = th.nn.ModuleList([
            # add layer norm
            th.nn.Linear(in_size, mid_size),
            th.nn.SiLU(),
            th.nn.Linear(mid_size, out_size)
        ])
        # initialize
        for proj in self.projectors:
            if isinstance(proj, th.nn.Linear):
                th.nn.init.xavier_uniform_(proj.weight)
                th.nn.init.zeros_(proj.bias)

    def forward(self, lower_in, upper_in, ):
        # ALso format lower_in
        lower_in = self.lower_projector(lower_in)
        upper_in = self.upper_projector(upper_in)
        x = th.cat([lower_in, upper_in], -1)
        for proj in self.projectors:
            x = proj(x)
        return x

class DinoSiglipMixer(th.nn.Module):
    def __init__(self, in_size=OUT_SIZE * 2, out_size=OUT_SIZE):
        super().__init__()
        self.dino_projector = HighLowMixer()
        self.siglip_projector = HighLowMixer()
        self.projectors = th.nn.Sequential(
            th.nn.SiLU(),
            th.nn.Linear(in_size, out_size),
        )
        # initialize
        for proj in self.projectors:
            if isinstance(proj, th.nn.Linear):
                th.nn.init.xavier_uniform_(proj.weight)
                th.nn.init.zeros_(proj.bias)

    
    def forward(self, dino_in, siglip_in):
        # ALso format lower_in
        lower, upper = th.chunk(dino_in, 2, -1)
        dino_out = self.dino_projector(lower, upper)
        lower, upper = th.chunk(siglip_in, 2, -1)
        siglip_out = self.siglip_projector(lower, upper)
        x = th.cat([dino_out, siglip_out], -1)
        for proj in self.projectors:
            x = proj(x)
        return x

class AnalogyEncoder(ModelMixin, ConfigMixin):
    @register_to_config
    def __init__(self, load_pretrained=False, 
                 dino_config_dict=None, siglip_config_dict=None):
        super().__init__()
        if load_pretrained:
            image_encoder_dino = AutoModel.from_pretrained('facebook/dinov2-large', torch_dtype=th.float16)
            image_encoder_siglip = SiglipVisionModel.from_pretrained("google/siglip-large-patch16-256", torch_dtype=th.float16, attn_implementation="sdpa")
        else:
            image_encoder_dino = AutoModel.from_config(Dinov2Config.from_dict(dino_config_dict))
            image_encoder_siglip = AutoModel.from_config(SiglipVisionConfig.from_dict(siglip_config_dict))
            
        image_encoder_dino.requires_grad_(False)
        image_encoder_dino = image_encoder_dino.to(memory_format=th.channels_last)

        image_encoder_siglip.requires_grad_(False)
        image_encoder_siglip = image_encoder_siglip.to(memory_format=th.channels_last)
        self.image_encoder_dino = image_encoder_dino
        self.image_encoder_siglip = image_encoder_siglip


    def dino_normalization(self, encoder_output):
        embeds = encoder_output.last_hidden_state
        embeds_pooled = embeds[:, 0:1]
        embeds = embeds / th.norm(embeds_pooled, dim=-1, keepdim=True)
        return embeds
    
    def siglip_normalization(self, encoder_output):
        embeds = th.cat ([encoder_output.pooler_output[:, None, :], encoder_output.last_hidden_state], dim=1)
        embeds_pooled = embeds[:, 0:1]
        embeds = embeds / th.norm(embeds_pooled, dim=-1, keepdim=True)
        return embeds
    
    def forward(self, dino_in, siglip_in):

        x_1 = self.image_encoder_dino(dino_in, output_hidden_states=True)
        x_1_first = x_1.hidden_states[0]
        x_1 = self.dino_normalization(x_1)
        x_2 = self.image_encoder_siglip(siglip_in, output_hidden_states=True)
        x_2_first = x_2.hidden_states[0]
        x_2_first_pool = th.mean(x_2_first, dim=1, keepdim=True)
        x_2_first = th.cat([x_2_first_pool, x_2_first], 1)
        x_2 = self.siglip_normalization(x_2)
        dino_embd = th.cat([x_1, x_1_first], -1)
        siglip_embd = th.cat([x_2, x_2_first], -1)
        return dino_embd, siglip_embd
    

class PatternAnalogyTrifuser(DiffusionPipeline):
    r"""
    This model inherits from [`DiffusionPipeline`]. Check the superclass documentation for the generic methods the
    library implements for all the pipelines (such as downloading or saving, running on a particular device, etc.)
    """

    model_cpu_offload_seq = "bert->unet->vqvae"

    analogy_input_processor: AnalogyInputProcessor
    analogy_encoder: AnalogyEncoder
    analogy_projector: AnalogyProjector
    unet: UNet2DConditionModel
    vae: AutoencoderKL
    scheduler: KarrasDiffusionSchedulers
    
    def __init__(self, 
        analogy_input_processor: AnalogyInputProcessor,
        analogy_projector: AnalogyProjector,
        analogy_encoder: AnalogyEncoder,
        unet: UNet2DConditionModel,
        vae: AutoencoderKL,
        scheduler: KarrasDiffusionSchedulers,):
        
        
        super().__init__()
        self.register_modules(
            analogy_input_processor=analogy_input_processor,
            analogy_encoder=analogy_encoder,
            analogy_projector=analogy_projector,
            unet=unet,
            vae=vae,
            scheduler=scheduler,
        )
        self.vae_scale_factor = 2 ** (len(self.vae.config.block_out_channels) - 1)
        self.image_processor = VaeImageProcessor(vae_scale_factor=self.vae_scale_factor)

    # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_image_variation.StableDiffusionImageVariationPipeline.check_inputs
    def check_inputs(self, analogy_prompt, negative_analogy_prompt, height, width, callback_steps):
        if (
            not isinstance(analogy_prompt, th.Tensor)
            and not isinstance(analogy_prompt, PIL.Image.Image)
            and not isinstance(analogy_prompt, list)
        ):
            raise ValueError(
                "`analogy_prompt` contents have to be of type `torch.Tensor` or `PIL.Image.Image` or `List[PIL.Image.Image]` but is"
                f" {type(analogy_prompt)}"
            )
        if not negative_analogy_prompt is None:
            if (
                not isinstance(negative_analogy_prompt, th.Tensor)
                and not isinstance(negative_analogy_prompt, PIL.Image.Image)
                and not isinstance(negative_analogy_prompt, list)
            ):
                raise ValueError(
                    "`negative_analogy_prompt` contents have to be of type `torch.Tensor` or `PIL.Image.Image` or `List[PIL.Image.Image]` but is"
                    f" {type(negative_analogy_prompt)}"
                )


        if height % 8 != 0 or width % 8 != 0:
            raise ValueError(f"`height` and `width` have to be divisible by 8 but are {height} and {width}.")

        if (callback_steps is None) or (
            callback_steps is not None and (not isinstance(callback_steps, int) or callback_steps <= 0)
        ):
            raise ValueError(
                f"`callback_steps` has to be a positive integer but is {callback_steps} of type"
                f" {type(callback_steps)}."
            )

    # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.prepare_latents
    def prepare_latents(self, batch_size, num_channels_latents, height, width, dtype, device, generator, latents=None):
        shape = (
            batch_size,
            num_channels_latents,
            int(height) // self.vae_scale_factor,
            int(width) // self.vae_scale_factor,
        )
        if isinstance(generator, list) and len(generator) != batch_size:
            raise ValueError(
                f"You have passed a list of generators of length {len(generator)}, but requested an effective batch"
                f" size of {batch_size}. Make sure the batch size matches the length of the generators."
            )

        if latents is None:
            latents = randn_tensor(shape, generator=generator, device=device, dtype=dtype)
        else:
            latents = latents.to(device)

        # scale the initial noise by the standard deviation required by the scheduler
        latents = latents * self.scheduler.init_noise_sigma
        return latents
    
    # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.prepare_extra_step_kwargs
    def prepare_extra_step_kwargs(self, generator, eta):
        # prepare extra kwargs for the scheduler step, since not all schedulers have the same signature
        # eta (η) is only used with the DDIMScheduler, it will be ignored for other schedulers.
        # eta corresponds to η in DDIM paper: https://arxiv.org/abs/2010.02502
        # and should be between [0, 1]

        accepts_eta = "eta" in set(inspect.signature(self.scheduler.step).parameters.keys())
        extra_step_kwargs = {}
        if accepts_eta:
            extra_step_kwargs["eta"] = eta

        # check if the scheduler accepts generator
        accepts_generator = "generator" in set(inspect.signature(self.scheduler.step).parameters.keys())
        if accepts_generator:
            extra_step_kwargs["generator"] = generator
        return extra_step_kwargs

    # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.prepare_latents
    def prepare_latents(self, batch_size, num_channels_latents, height, width, dtype, device, generator, latents=None):
        shape = (
            batch_size,
            num_channels_latents,
            int(height) // self.vae_scale_factor,
            int(width) // self.vae_scale_factor,
        )
        if isinstance(generator, list) and len(generator) != batch_size:
            raise ValueError(
                f"You have passed a list of generators of length {len(generator)}, but requested an effective batch"
                f" size of {batch_size}. Make sure the batch size matches the length of the generators."
            )

        if latents is None:
            latents = randn_tensor(shape, generator=generator, device=device, dtype=dtype)
        else:
            latents = latents.to(device)

        # scale the initial noise by the standard deviation required by the scheduler
        latents = latents * self.scheduler.init_noise_sigma
        return latents
    
    def _encode_prompt(self, analogy_prompt, device, num_images_per_prompt, do_classifier_free_guidance, negative_prompt):
        r"""
        Encodes the prompt into text encoder hidden states.

        Args:
            prompt (`str` or `List[str]`):
                prompt to be encoded
            device: (`torch.device`):
                torch device
            num_images_per_prompt (`int`):
                number of images that should be generated per prompt
            do_classifier_free_guidance (`bool`):
                whether to use classifier free guidance or not
            negative_prompt (`str` or `List[str]`):
                The prompt or prompts not to guide the image generation. Ignored when not using guidance (i.e., ignored
                if `guidance_scale` is less than `1`).
        """
        weight_dtype = self.unet.dtype
        dino_input, siglip_input = self.analogy_input_processor(analogy_prompt)
        dino_input = dino_input.to(device=device).to(dtype=weight_dtype)
        siglip_input = siglip_input.to(device=device).to(dtype=weight_dtype)
        batch_size = dino_input.shape[1]
        dino_input_reshaped = einops.rearrange(dino_input, "k b c h w -> (k b) c h w")
        siglip_input_reshaped = einops.rearrange(siglip_input, "k b c h w -> (k b) c h w")
        dino_enc, siglip_enc = self.analogy_encoder(dino_input_reshaped, siglip_input_reshaped)
        image_embeddings = self.analogy_projector(dino_enc, siglip_enc, batch_size)
        # Check size here.
        
        bs_embed, seq_len, _ = image_embeddings.shape
        image_embeddings = image_embeddings.repeat(num_images_per_prompt, 1, 1)
        # get unconditional embeddings for classifier free guidance
        if do_classifier_free_guidance:
            uncond_images: List[str]
            if negative_prompt is None:
                uncond_images = [np.zeros((512, 512, 3)) + 0.5] * batch_size
            elif type(negative_prompt) is not type(analogy_prompt):
                raise TypeError(
                    f"`negative_prompt` should be the same type to `prompt`, but got {type(analogy_prompt)} !="
                    f" {type(negative_prompt)}."
                )
            elif isinstance(negative_prompt, PIL.Image.Image):
                uncond_images = [negative_prompt]
            elif batch_size != len(negative_prompt):
                raise ValueError(
                    f"`negative_prompt`: {negative_prompt} has batch size {len(negative_prompt)}, but `prompt`:"
                    f" {analogy_prompt} has batch size {batch_size}. Please make sure that passed `negative_prompt` matches"
                    " the batch size of `prompt`."
                )
            else:
                uncond_images = negative_prompt
            dino_neg, siglip_neg = self.analogy_input_processor.get_negative(dino_input, siglip_input)
            
            dino_neg = dino_neg.to(device=device).to(dtype=weight_dtype)
            siglip_neg = siglip_neg.to(device=device).to(dtype=weight_dtype)
            dino_neg_reshaped = einops.rearrange(dino_neg, "k b c h w -> (k b) c h w")
            siglip_neg_reshaped = einops.rearrange(siglip_neg, "k b c h w -> (k b) c h w")
            dino_neg_enc, siglip_neg_enc = self.analogy_encoder(dino_neg_reshaped, siglip_neg_reshaped)
            negative_prompt_embeds = self.analogy_projector(dino_neg_enc, siglip_neg_enc, batch_size)
                
            negative_prompt_embeds = negative_prompt_embeds.repeat(num_images_per_prompt, 1, 1)
            image_embeddings = th.cat([negative_prompt_embeds, image_embeddings])


        return image_embeddings
    
    @th.no_grad()
    def __call__(
        self,
        analogy_prompt: Union[str, List[str]] = None,
        num_inference_steps: int = 50,
        guidance_scale: float = 7.5,
        height: Optional[int] = None,
        width: Optional[int] = None,
        negative_analogy_prompt: Optional[Union[str, List[str]]] = None,
        num_images_per_prompt: Optional[int] = 1,
        eta: float = 0.0,
        generator: Optional[Union[th.Generator, List[th.Generator]]] = None,
        latents: Optional[th.FloatTensor] = None,
        output_type: Optional[str] = "pil",
        return_dict: bool = True,
        callback: Optional[Callable[[int, int, th.Tensor], None]] = None,
        callback_steps: int = 1,
        start_step: int = 0,
    ):
        r"""
        The call function to the pipeline for generation.

        Args:
            image (`PIL.Image.Image`, `List[PIL.Image.Image]` or `torch.Tensor`):
                The image prompt or prompts to guide the image generation.
            height (`int`, *optional*, defaults to `self.image_unet.config.sample_size * self.vae_scale_factor`):
                The height in pixels of the generated image.
            width (`int`, *optional*, defaults to `self.image_unet.config.sample_size * self.vae_scale_factor`):
                The width in pixels of the generated image.
            num_inference_steps (`int`, *optional*, defaults to 50):
                The number of denoising steps. More denoising steps usually lead to a higher quality image at the
                expense of slower inference.
            guidance_scale (`float`, *optional*, defaults to 7.5):
                A higher guidance scale value encourages the model to generate images closely linked to the text
                `prompt` at the expense of lower image quality. Guidance scale is enabled when `guidance_scale > 1`.
            negative_prompt (`str` or `List[str]`, *optional*):
                The prompt or prompts to guide what to not include in image generation. If not defined, you need to
                pass `negative_prompt_embeds` instead. Ignored when not using guidance (`guidance_scale < 1`).
            num_images_per_prompt (`int`, *optional*, defaults to 1):
                The number of images to generate per prompt.
            eta (`float`, *optional*, defaults to 0.0):
                Corresponds to parameter eta (η) from the [DDIM](https://arxiv.org/abs/2010.02502) paper. Only applies
                to the [`~schedulers.DDIMScheduler`], and is ignored in other schedulers.
            generator (`torch.Generator`, *optional*):
                A [`torch.Generator`](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make
                generation deterministic.
            latents (`torch.Tensor`, *optional*):
                Pre-generated noisy latents sampled from a Gaussian distribution, to be used as inputs for image
                generation. Can be used to tweak the same generation with different prompts. If not provided, a latents
                tensor is generated by sampling using the supplied random `generator`.
            output_type (`str`, *optional*, defaults to `"pil"`):
                The output format of the generated image. Choose between `PIL.Image` or `np.array`.
            return_dict (`bool`, *optional*, defaults to `True`):
                Whether or not to return a [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] instead of a
                plain tuple.
            callback (`Callable`, *optional*):
                A function that calls every `callback_steps` steps during inference. The function is called with the
                following arguments: `callback(step: int, timestep: int, latents: torch.Tensor)`.
            callback_steps (`int`, *optional*, defaults to 1):
                The frequency at which the `callback` function is called. If not specified, the callback is called at
                every step.

        Examples:

        ```py
        >>> from diffusers import VersatileDiffusionImageVariationPipeline
        >>> import torch
        >>> import requests
        >>> from io import BytesIO
        >>> from PIL import Image

        >>> # let's download an initial image
        >>> url = "https://huggingface.co/datasets/diffusers/images/resolve/main/benz.jpg"

        >>> response = requests.get(url)
        >>> image = Image.open(BytesIO(response.content)).convert("RGB")

        >>> pipe = VersatileDiffusionImageVariationPipeline.from_pretrained(
        ...     "shi-labs/versatile-diffusion", torch_dtype=torch.float16
        ... )
        >>> pipe = pipe.to("cuda")

        >>> generator = torch.Generator(device="cuda").manual_seed(0)
        >>> image = pipe(image, generator=generator).images[0]
        >>> image.save("./car_variation.png")
        ```

        Returns:
            [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] or `tuple`:
                If `return_dict` is `True`, [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] is returned,
                otherwise a `tuple` is returned where the first element is a list with the generated images.
        """

        # 1. Check inputs. Raise error if not correct
        height = height or self.unet.config.sample_size * self.vae_scale_factor
        width = width or self.unet.config.sample_size * self.vae_scale_factor

        # 1. Check inputs. Raise error if not correct
        self.check_inputs(analogy_prompt, negative_analogy_prompt, height, width, callback_steps)
        
        # 2. Define call parameters
        if isinstance(analogy_prompt, list):
            batch_size = len(analogy_prompt)
        elif isinstance(analogy_prompt, tuple):
            batch_size = 1
        else:
            raise ValueError(
                f"`analogy_prompt` has to be a list of images or a tuple of images but is of type {type(analogy_prompt)}"
            )
        device = self._execution_device
        # here `guidance_scale` is defined analog to the guidance weight `w` of equation (2)
        # of the Imagen paper: https://arxiv.org/pdf/2205.11487.pdf . `guidance_scale = 1`
        # corresponds to doing no classifier free guidance.
        do_classifier_free_guidance = guidance_scale > 1.0

        # 3. Encode input prompt
        analogy_embeddings = self._encode_prompt(
            analogy_prompt, device, num_images_per_prompt, do_classifier_free_guidance, negative_analogy_prompt
        )

        # 4. Prepare timesteps
        self.scheduler.set_timesteps(num_inference_steps, device=device)
        
        timesteps = self.scheduler.timesteps
        # Now this should be from start step onwards
        timesteps = timesteps[start_step:]
        # 5. Prepare latent variables
        num_channels_latents = self.unet.config.in_channels
        latents = self.prepare_latents(
            batch_size * num_images_per_prompt,
            num_channels_latents,
            height,
            width,
            analogy_embeddings.dtype,
            device,
            generator,
            latents,
        )

        # 6. Prepare extra step kwargs.
        extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta)

        # 7. Denoising loop
        for i, t in enumerate(self.progress_bar(timesteps)):
            # expand the latents if we are doing classifier free guidance
            latent_model_input = th.cat([latents] * 2) if do_classifier_free_guidance else latents
            latent_model_input = self.scheduler.scale_model_input(latent_model_input, t)

            # predict the noise residual
            noise_pred = self.unet(latent_model_input, t, encoder_hidden_states=analogy_embeddings).sample

            # perform guidance
            if do_classifier_free_guidance:
                noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)
                noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)

            # compute the previous noisy sample x_t -> x_t-1
            latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs).prev_sample

            # call the callback, if provided
            if callback is not None and i % callback_steps == 0:
                step_idx = i // getattr(self.scheduler, "order", 1)
                callback(step_idx, t, latents)

        if not output_type == "latent":
            image = self.vae.decode(latents / self.vae.config.scaling_factor, return_dict=False)[0]
        else:
            image = latents

        image = self.image_processor.postprocess(image, output_type=output_type)

        if not return_dict:
            return (image,)

        return ImagePipelineOutput(images=image)