pipeline.py

import random
from tqdm import tqdm
from typing import Callable, Dict, List, Optional

import torch
from diffusers import DiffusionPipeline
from diffusers.configuration_utils import ConfigMixin


class SuperDiffPipeline(DiffusionPipeline, ConfigMixin):
    """SuperDiffPipeline."""

    def __init__(self, unet: Callable, vae: Callable, text_encoder: Callable, scheduler: Callable, tokenizer: Callable) -> None:
        """__init__.

        Parameters
        ----------
        unet : Callable
            unet
        vae : Callable
            vae
        text_encoder : Callable
            text_encoder
        scheduler : Callable
            scheduler
        tokenizer : Callable
            tokenizer
        kwargs :
            kwargs

        Returns
        -------
        None

        """
        super().__init__()

        device = "cuda" if torch.cuda.is_available() else "cpu"

        vae.to(device)
        unet.to(device)
        text_encoder.to(device)
        self.register_modules(unet=unet, 
                              scheduler=scheduler, 
                              vae=vae, 
                              text_encoder=text_encoder,
                              tokenizer=tokenizer,)

    @torch.no_grad
    def get_batch(self, latents: Callable, nrow: int, ncol: int) -> Callable:
        """get_batch.

        Parameters
        ----------
        latents : Callable
            latents
        nrow : int
            nrow
        ncol : int
            ncol

        Returns
        -------
        Callable

        """
        image = self.vae.decode(
            latents / self.vae.config.scaling_factor, return_dict=False
        )[0]
        image = (image / 2 + 0.5).clamp(0, 1).squeeze()
        if len(image.shape) < 4:
            image = image.unsqueeze(0)
        image = (image.permute(0, 2, 3, 1) * 255).to(torch.uint8)
        return image

    @torch.no_grad
    def get_text_embedding(self, prompt: str) -> Callable:
        """get_text_embedding.

        Parameters
        ----------
        prompt : str
            prompt

        Returns
        -------
        Callable

        """
        text_input = self.tokenizer(
            prompt,
            padding="max_length",
            max_length=self.tokenizer.model_max_length,
            truncation=True,
            return_tensors="pt",
        )
        return self.text_encoder(text_input.input_ids.to(self.device))[0]

    @torch.no_grad
    def get_vel(self, t: float, sigma: float, latents: Callable, embeddings: Callable):
        """get_vel.

        Parameters
        ----------
        t : float
            t
        sigma : float
            sigma
        latents : Callable
            latents
        embeddings : Callable
            embeddings
        """
        def v(_x, _e): return self.unet(
            _x / ((sigma**2 + 1) ** 0.5), t, encoder_hidden_states=_e
        ).sample
        embeds = torch.cat(embeddings)
        latent_input = latents
        vel = v(latent_input, embeds)
        return vel

    def preprocess(
        self,
        prompt_1: str,
        prompt_2: str,
        seed: int = None,
        num_inference_steps: int = 1000,
        batch_size: int = 1,
        lift: int = 0.0,
        height: int = 512,
        width: int = 512,
        guidance_scale: int = 7.5,
    ) -> Callable:
        """preprocess.

        Parameters
        ----------
        prompt_1 : str
            prompt_1
        prompt_2 : str
            prompt_2
        seed : int
            seed
        num_inference_steps : int
            num_inference_steps
        batch_size : int
            batch_size
        lift : int
            lift
        height : int
            height
        width : int
            width
        guidance_scale : int
            guidance_scale

        Returns
        -------
        Callable

        """
        # Tokenize the input
        self.batch_size = batch_size
        self.num_inference_steps = num_inference_steps
        self.guidance_scale = guidance_scale
        self.lift = lift
        self.seed = seed
        if self.seed is None:
            self.seed = random.randint(0, 2**32 - 1)
        obj_prompt = [prompt_1]
        bg_prompt = [prompt_2]
        obj_embeddings = self.get_text_embedding(obj_prompt * batch_size)
        bg_embeddings = self.get_text_embedding(bg_prompt * batch_size)

        uncond_embeddings = self.get_text_embedding([""] * batch_size)

        generator = torch.cuda.manual_seed(
            self.seed
        )  # Seed generator to create the initial latent noise
        latents = torch.randn(
            (batch_size, self.unet.config.in_channels, height // 8, width // 8),
            generator=generator,
            device=self.device,
        )

        latents_og = latents.clone().detach()
        latents_uncond_og = latents.clone().detach()

        self.scheduler.set_timesteps(num_inference_steps)
        latents = latents * self.scheduler.init_noise_sigma

        latents_uncond = latents.clone().detach()
        return {
            "latents": latents,
            "obj_embeddings": obj_embeddings,
            "uncond_embeddings": uncond_embeddings,
            "bg_embeddings": bg_embeddings,
        }

    def _forward(self, model_inputs: Dict) -> Callable:
        """_forward.

        Parameters
        ----------
        model_inputs : Dict
            model_inputs

        Returns
        -------
        Callable

        """
        latents = model_inputs["latents"]
        obj_embeddings = model_inputs["obj_embeddings"]
        uncond_embeddings = model_inputs["uncond_embeddings"]
        bg_embeddings = model_inputs["bg_embeddings"]

        kappa = 0.5 * torch.ones(
            (self.num_inference_steps + 1, self.batch_size), device=self.device
        )
        ll_obj = torch.ones(
            (self.num_inference_steps + 1, self.batch_size), device=self.device
        )
        ll_bg = torch.ones(
            (self.num_inference_steps + 1, self.batch_size), device=self.device
        )
        ll_uncond = torch.ones(
            (self.num_inference_steps + 1, self.batch_size), device=self.device
        )
        with torch.no_grad():
            for i, t in tqdm(enumerate(self.scheduler.timesteps)):
                dsigma = self.scheduler.sigmas[i +
                                               1] - self.scheduler.sigmas[i]
                sigma = self.scheduler.sigmas[i]
                vel_obj = self.get_vel(t, sigma, latents, [obj_embeddings])
                vel_uncond = self.get_vel(
                    t, sigma, latents, [uncond_embeddings])

                vel_bg = self.get_vel(t, sigma, latents, [bg_embeddings])
                noise = torch.sqrt(2 * torch.abs(dsigma) * sigma) * torch.randn_like(
                    latents
                )

                dx_ind = (
                    2
                    * dsigma
                    * (vel_uncond + self.guidance_scale * (vel_bg - vel_uncond))
                    + noise
                )
                kappa[i + 1] = (
                    (torch.abs(dsigma) * (vel_bg - vel_obj) * (vel_bg + vel_obj)).sum(
                        (1, 2, 3)
                    )
                    - (dx_ind * ((vel_obj - vel_bg))).sum((1, 2, 3))
                    + sigma * self.lift / self.num_inference_steps
                )
                kappa[i + 1] /= (
                    2
                    * dsigma
                    * self.guidance_scale
                    * ((vel_obj - vel_bg) ** 2).sum((1, 2, 3))
                )

                vf = vel_uncond + self.guidance_scale * (
                    (vel_bg - vel_uncond)
                    + kappa[i + 1][:, None, None, None] * (vel_obj - vel_bg)
                )
                dx = 2 * dsigma * vf + noise
                latents += dx

                ll_obj[i + 1] = ll_obj[i] + (
                    -torch.abs(dsigma) / sigma * (vel_obj) ** 2
                    - (dx * (vel_obj / sigma))
                ).sum((1, 2, 3))
                ll_bg[i + 1] = ll_bg[i] + (
                    -torch.abs(dsigma) / sigma * (vel_bg) ** 2 -
                    (dx * (vel_bg / sigma))
                ).sum((1, 2, 3))

        return latents

    def postprocess(self, latents: Callable) -> Callable:
        """postprocess.

        Parameters
        ----------
        latents : Callable
            latents

        Returns
        -------
        Callable

        """
        image = self.get_batch(latents, 1, self.batch_size)
        # Ensure the shape is (height, width, 3)
        assert image.shape[-1] == 3  # Handle grayscale or invalid shapes

        # Convert to uint8 if not already
        image = image.to(torch.uint8)  # Ensure it's uint8 for PIL

        return image

    def __call__(
        self,
        prompt_1: str,
        prompt_2: str,
        seed: int = None,
        num_inference_steps: int = 1000,
        batch_size: int = 1,
        lift: int = 0.0,
        height: int = 512,
        width: int = 512,
        guidance_scale: int = 7.5,
    ) -> Callable:
        """__call__.

        Parameters
        ----------
        prompt_1 : str
            prompt_1
        prompt_2 : str
            prompt_2
        seed : int
            seed
        num_inference_steps : int
            num_inference_steps
        batch_size : int
            batch_size
        lift : int
            lift
        height : int
            height
        width : int
            width
        guidance_scale : int
            guidance_scale

        Returns
        -------
        Callable

        """
        # Preprocess inputs
        model_inputs = self.preprocess(
            prompt_1,
            prompt_2,
            seed,
            num_inference_steps,
            batch_size,
            lift,
            height,
            width,
            guidance_scale,
        )

        # Forward pass through the pipeline
        latents = self._forward(model_inputs)

        # Postprocess to generate the final output
        images = self.postprocess(latents)
        return images