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

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 analogy_encoder import AnalogyEncoder
from analogy_projector import AnalogyProjector
from analogy_input_processor import AnalogyInputProcessor

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:
            analogy_prompt (`List[Tuple[PIL.Image.Image]]'):
                The analogy sequence A, A*, B which is our model's prompt for generating B* the analogical pattern satisfying A:A*::B:B*.
            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
        import requests
        import torch as th
        from PIL import Image
        from io import BytesIO
        import matplotlib.pyplot as plt
        from PIL import Image, ImageOps
        from diffusers import DiffusionPipeline

        SEED = 1729
        DEVICE = th.device("cuda")
        DTYPE = th.float16
        FIG_K = 3
        EXAMPLE_ID = 0

        # Now we need to do the trick
        pretrained_path = "bardofcodes/pattern_analogies"
        new_pipe = DiffusionPipeline.from_pretrained(
            pretrained_path, 
            custom_pipeline=pretrained_path, 
            trust_remote_code=True
        )

        img_urls = [
            f"https://huggingface.co./bardofcodes/pattern_analogies/resolve/main/examples/{EXAMPLE_ID}_a.png",
            f"https://huggingface.co./bardofcodes/pattern_analogies/resolve/main/examples/{EXAMPLE_ID}_a_star.png",
            f"https://huggingface.co./bardofcodes/pattern_analogies/resolve/main/examples/{EXAMPLE_ID}_b.png",
        ]
        images = []
        for url in img_urls:
            response = requests.get(url)
            image = Image.open(BytesIO(response.content)).convert("RGB")
            images.append(image)

        pipe_input = [tuple(images)]

        pipe = new_pipe.to(DEVICE, DTYPE)
        var_images = pipe(pipe_input, num_inference_steps=50, num_images_per_prompt=3,).images

        plt.figure(figsize=(3*FIG_K, 2*FIG_K))
        plt.axis('off')
        plt.legend(framealpha=1)
        plt.rcParams['legend.fontsize'] = 'large'
        for i in range(6):
            if i == 0:
                plt.subplot(2, 3, i+1)
                val_image = img1
                label_str = "A"
            elif i == 1:
                plt.subplot(2, 3, i+1)
                val_image = alt_img1
                label_str = "A*"
            elif i == 2:
                plt.subplot(2, 3, i+1)
                val_image = img2
                label_str = "Target"
            else:
                plt.subplot(2, 3,i + 1)
                val_image = var_images[i-3]
                label_str = f"Variation {i-2}"
            
            val_image = ImageOps.expand(val_image,border=2,fill='black')
            plt.imshow(val_image)
            plt.scatter([], [], c="r", label=label_str)
            plt.legend(loc="lower right")
            plt.axis('off')
        plt.subplots_adjust(wspace=0.01, hspace=0.01)
        ```

        Returns:
            [`~ImagePipelineOutput`] or `tuple`
            The generated image(s) as a [`~ImagePipelineOutput`] or a tuple of 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)