pipeline_sd.py

import copy
import math
from typing import Any, Dict, List, Optional, Tuple, Union

import torch
import torch.nn.functional as F
import utils
from accelerate import Accelerator
from diffusers import StableDiffusionPipeline
from diffusers.image_processor import PipelineImageInput
from losses import *
from tqdm import tqdm


class ADPipeline(StableDiffusionPipeline):
    def freeze(self):
        self.vae.requires_grad_(False)
        self.unet.requires_grad_(False)
        self.text_encoder.requires_grad_(False)
        self.classifier.requires_grad_(False)

    @torch.no_grad()
    def image2latent(self, image):
        dtype = next(self.vae.parameters()).dtype
        device = self._execution_device
        image = image.to(device=device, dtype=dtype) * 2.0 - 1.0
        latent = self.vae.encode(image)["latent_dist"].mean
        latent = latent * self.vae.config.scaling_factor
        return latent

    @torch.no_grad()
    def latent2image(self, latent):
        dtype = next(self.vae.parameters()).dtype
        device = self._execution_device
        latent = latent.to(device=device, dtype=dtype)
        latent = latent / self.vae.config.scaling_factor
        image = self.vae.decode(latent)[0]
        return (image * 0.5 + 0.5).clamp(0, 1)

    def init(self, enable_gradient_checkpoint):
        self.freeze()
        weight_dtype = torch.float32
        if self.accelerator.mixed_precision == "fp16":
            weight_dtype = torch.float16
        elif self.accelerator.mixed_precision == "bf16":
            weight_dtype = torch.bfloat16

        # Move unet, vae and text_encoder to device and cast to weight_dtype
        self.unet.to(self.accelerator.device, dtype=weight_dtype)
        self.vae.to(self.accelerator.device, dtype=weight_dtype)
        self.text_encoder.to(self.accelerator.device, dtype=weight_dtype)
        self.classifier.to(self.accelerator.device, dtype=weight_dtype)
        self.classifier = self.accelerator.prepare(self.classifier)
        if enable_gradient_checkpoint:
            self.classifier.enable_gradient_checkpointing()

    def sample(
        self,
        lr=0.05,
        iters=1,
        attn_scale=1,
        adain=False,
        weight=0.25,
        controller=None,
        style_image=None,
        content_image=None,
        mixed_precision="no",
        start_time=999,
        enable_gradient_checkpoint=False,
        prompt: Union[str, List[str]] = None,
        height: Optional[int] = None,
        width: Optional[int] = None,
        num_inference_steps: int = 50,
        guidance_scale: float = 7.5,
        negative_prompt: Optional[Union[str, List[str]]] = None,
        num_images_per_prompt: Optional[int] = 1,
        generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
        latents: Optional[torch.Tensor] = None,
        prompt_embeds: Optional[torch.Tensor] = None,
        negative_prompt_embeds: Optional[torch.Tensor] = None,
        ip_adapter_image: Optional[PipelineImageInput] = None,
        ip_adapter_image_embeds: Optional[List[torch.Tensor]] = None,
        cross_attention_kwargs: Optional[Dict[str, Any]] = None,
        guidance_rescale: float = 0.0,
        clip_skip: Optional[int] = None,
        **kwargs,
    ):
        # 0. Default height and width to unet
        height = height or self.unet.config.sample_size * self.vae_scale_factor
        width = width or self.unet.config.sample_size * self.vae_scale_factor
        self._guidance_scale = guidance_scale
        self._guidance_rescale = guidance_rescale
        self._clip_skip = clip_skip
        self._cross_attention_kwargs = cross_attention_kwargs
        self._interrupt = False

        self.accelerator = Accelerator(
            mixed_precision=mixed_precision, gradient_accumulation_steps=1
        )
        self.init(enable_gradient_checkpoint)

        # 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
        lora_scale = (
            self.cross_attention_kwargs.get("scale", None)
            if self.cross_attention_kwargs is not None
            else None
        )
        do_cfg = guidance_scale > 1.0

        prompt_embeds, negative_prompt_embeds = self.encode_prompt(
            prompt,
            device,
            num_images_per_prompt,
            do_cfg,
            negative_prompt,
            prompt_embeds=prompt_embeds,
            negative_prompt_embeds=negative_prompt_embeds,
            lora_scale=lora_scale,
            clip_skip=self.clip_skip,
        )

        # For classifier free guidance, we need to do two forward passes.
        # Here we concatenate the unconditional and text embeddings into a single batch
        # to avoid doing two forward passes
        if do_cfg:
            prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds])

        if ip_adapter_image is not None or ip_adapter_image_embeds is not None:
            image_embeds = self.prepare_ip_adapter_image_embeds(
                ip_adapter_image,
                ip_adapter_image_embeds,
                device,
                batch_size * num_images_per_prompt,
                do_cfg,
            )

        # 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,
            prompt_embeds.dtype,
            device,
            generator,
            latents,
        )

        # 6.1 Add image embeds for IP-Adapter
        added_cond_kwargs = (
            {"image_embeds": image_embeds}
            if (ip_adapter_image is not None or ip_adapter_image_embeds is not None)
            else None
        )

        # 6.2 Optionally get Guidance Scale Embedding
        timestep_cond = None
        if self.unet.config.time_cond_proj_dim is not None:
            guidance_scale_tensor = torch.tensor(self.guidance_scale - 1).repeat(
                batch_size * num_images_per_prompt
            )
            timestep_cond = self.get_guidance_scale_embedding(
                guidance_scale_tensor, embedding_dim=self.unet.config.time_cond_proj_dim
            ).to(device=device, dtype=latents.dtype)

        self.scheduler.set_timesteps(num_inference_steps)
        timesteps = self.scheduler.timesteps
        self.style_latent = self.image2latent(style_image)
        if content_image is not None:
            self.content_latent = self.image2latent(content_image)
        else:
            self.content_latent = None
        null_embeds = self.encode_prompt("", device, 1, False)[0]
        self.null_embeds = null_embeds
        self.null_embeds_for_latents = torch.cat([null_embeds] * latents.shape[0])
        self.null_embeds_for_style = torch.cat(
            [null_embeds] * self.style_latent.shape[0]
        )
        
        self.adain = adain
        self.attn_scale = attn_scale
        self.cache = utils.DataCache()
        self.controller = controller
        utils.register_attn_control(
            self.classifier, controller=self.controller, cache=self.cache
        )
        print("Total self attention layers of Unet: ", controller.num_self_layers)
        print("Self attention layers for AD: ", controller.self_layers)

        pbar = tqdm(timesteps, desc="Sample")
        for i, t in enumerate(pbar):
            with torch.no_grad():
                # expand the latents if we are doing classifier free guidance
                latent_model_input = torch.cat([latents] * 2) if do_cfg 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=prompt_embeds,
                    timestep_cond=timestep_cond,
                    cross_attention_kwargs=self.cross_attention_kwargs,
                    added_cond_kwargs=added_cond_kwargs,
                    return_dict=False,
                )[0]

                # perform guidance
                if do_cfg:
                    noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)
                    noise_pred = noise_pred_uncond + self.guidance_scale * (
                        noise_pred_text - noise_pred_uncond
                    )
                latents = self.scheduler.step(
                    noise_pred, t, latents, return_dict=False
                )[0]
            if iters > 0 and t < start_time:
                latents = self.AD(latents, t, lr, iters, pbar, weight)
                
        images = self.latent2image(latents)
        # Offload all models
        self.maybe_free_model_hooks()
        return images

    def optimize(
        self,
        latents=None,
        attn_scale=1.0,
        lr=0.05,
        iters=1,
        weight=0,
        width=512,
        height=512,
        batch_size=1,
        controller=None,
        style_image=None,
        content_image=None,
        mixed_precision="no",
        num_inference_steps=50,
        enable_gradient_checkpoint=False,
        source_mask=None,
        target_mask=None,
    ):
        height = height // self.vae_scale_factor
        width = width // self.vae_scale_factor

        self.accelerator = Accelerator(
            mixed_precision=mixed_precision, gradient_accumulation_steps=1
        )
        self.init(enable_gradient_checkpoint)

        style_latent = self.image2latent(style_image)
        latents = torch.randn((batch_size, 4, height, width), device=self.device)
        null_embeds = self.encode_prompt("", self.device, 1, False)[0]
        null_embeds_for_latents = null_embeds.repeat(latents.shape[0], 1, 1)
        null_embeds_for_style = null_embeds.repeat(style_latent.shape[0], 1, 1)

        if content_image is not None:
            content_latent = self.image2latent(content_image)
            latents = torch.cat([content_latent.clone()] * batch_size)
            null_embeds_for_content = null_embeds.repeat(content_latent.shape[0], 1, 1)

        self.cache = utils.DataCache()
        self.controller = controller
        utils.register_attn_control(
            self.classifier, controller=self.controller, cache=self.cache
        )
        print("Total self attention layers of Unet: ", controller.num_self_layers)
        print("Self attention layers for AD: ", controller.self_layers)

        self.scheduler.set_timesteps(num_inference_steps)
        timesteps = self.scheduler.timesteps
        latents = latents.detach().float()
        optimizer = torch.optim.Adam([latents.requires_grad_()], lr=lr)
        optimizer = self.accelerator.prepare(optimizer)
        pbar = tqdm(timesteps, desc="Optimize")
        for i, t in enumerate(pbar):
            # t = torch.tensor([1], device=self.device)
            with torch.no_grad():
                qs_list, ks_list, vs_list, s_out_list = self.extract_feature(
                    style_latent,
                    t,
                    null_embeds_for_style,
                )
                if content_image is not None:
                    qc_list, kc_list, vc_list, c_out_list = self.extract_feature(
                        content_latent,
                        t,
                        null_embeds_for_content,
                    )
            for j in range(iters):
                style_loss = 0
                content_loss = 0
                optimizer.zero_grad()
                q_list, k_list, v_list, self_out_list = self.extract_feature(
                    latents,
                    t,
                    null_embeds_for_latents,
                )
                style_loss = ad_loss(q_list, ks_list, vs_list, self_out_list, scale=attn_scale, source_mask=source_mask, target_mask=target_mask)
                if content_image is not None:
                    content_loss = q_loss(q_list, qc_list)
                    # content_loss = qk_loss(q_list, k_list, qc_list, kc_list)
                    # content_loss = qkv_loss(q_list, k_list, vc_list, c_out_list)
                loss = style_loss + content_loss * weight
                self.accelerator.backward(loss)
                optimizer.step()
                pbar.set_postfix(loss=loss.item(), time=t.item(), iter=j)
        images = self.latent2image(latents)
        # Offload all models
        self.maybe_free_model_hooks()
        return images

    def panorama(
        self,
        lr=0.05,
        iters=1,
        attn_scale=1,
        adain=False,
        controller=None,
        style_image=None,
        mixed_precision="no",
        enable_gradient_checkpoint=False,
        prompt: Union[str, List[str]] = None,
        height: Optional[int] = None,
        width: Optional[int] = None,
        num_inference_steps: int = 50,
        guidance_scale: float = 1,
        stride=8,
        view_batch_size: int = 16,
        negative_prompt: 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.Tensor] = None,
        prompt_embeds: Optional[torch.Tensor] = None,
        negative_prompt_embeds: Optional[torch.Tensor] = None,
        ip_adapter_image: Optional[PipelineImageInput] = None,
        ip_adapter_image_embeds: Optional[List[torch.Tensor]] = None,
        cross_attention_kwargs: Optional[Dict[str, Any]] = None,
        guidance_rescale: float = 0.0,
        clip_skip: Optional[int] = None,
        **kwargs,
    ):

        # 0. Default height and width to unet
        height = height or self.unet.config.sample_size * self.vae_scale_factor
        width = width or self.unet.config.sample_size * self.vae_scale_factor

        self._guidance_scale = guidance_scale
        self._guidance_rescale = guidance_rescale
        self._clip_skip = clip_skip
        self._cross_attention_kwargs = cross_attention_kwargs
        self._interrupt = False

        self.accelerator = Accelerator(
            mixed_precision=mixed_precision, gradient_accumulation_steps=1
        )
        self.init(enable_gradient_checkpoint)

        # 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
        # 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_cfg = guidance_scale > 1.0

        if ip_adapter_image is not None or ip_adapter_image_embeds is not None:
            image_embeds = self.prepare_ip_adapter_image_embeds(
                ip_adapter_image,
                ip_adapter_image_embeds,
                device,
                batch_size * num_images_per_prompt,
                self.do_classifier_free_guidance,
            )

        # 3. Encode input prompt
        text_encoder_lora_scale = (
            cross_attention_kwargs.get("scale", None)
            if cross_attention_kwargs is not None
            else None
        )
        prompt_embeds, negative_prompt_embeds = self.encode_prompt(
            prompt,
            device,
            num_images_per_prompt,
            do_cfg,
            negative_prompt,
            prompt_embeds=prompt_embeds,
            negative_prompt_embeds=negative_prompt_embeds,
            lora_scale=text_encoder_lora_scale,
            clip_skip=clip_skip,
        )
        # For classifier free guidance, we need to do two forward passes.
        # Here we concatenate the unconditional and text embeddings into a single batch
        # to avoid doing two forward passes
        if do_cfg:
            prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds])

        # 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,
            prompt_embeds.dtype,
            device,
            generator,
            latents,
        )

        # 6. Define panorama grid and initialize views for synthesis.
        # prepare batch grid
        views = self.get_views_(height, width, window_size=64, stride=stride)
        views_batch = [
            views[i : i + view_batch_size]
            for i in range(0, len(views), view_batch_size)
        ]
        print(len(views), len(views_batch), views_batch)
        self.scheduler.set_timesteps(num_inference_steps)
        views_scheduler_status = [copy.deepcopy(self.scheduler.__dict__)] * len(
            views_batch
        )
        count = torch.zeros_like(latents)
        value = torch.zeros_like(latents)

        # 7. Prepare extra step kwargs. TODO: Logic should ideally just be moved out of the pipeline
        extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta)

        # 7.1 Add image embeds for IP-Adapter
        added_cond_kwargs = (
            {"image_embeds": image_embeds}
            if ip_adapter_image is not None or ip_adapter_image_embeds is not None
            else None
        )

        # 7.2 Optionally get Guidance Scale Embedding
        timestep_cond = None
        if self.unet.config.time_cond_proj_dim is not None:
            guidance_scale_tensor = torch.tensor(self.guidance_scale - 1).repeat(
                batch_size * num_images_per_prompt
            )
            timestep_cond = self.get_guidance_scale_embedding(
                guidance_scale_tensor, embedding_dim=self.unet.config.time_cond_proj_dim
            ).to(device=device, dtype=latents.dtype)

        # 8. Denoising loop
        # Each denoising step also includes refinement of the latents with respect to the
        # views.

        timesteps = self.scheduler.timesteps
        self.style_latent = self.image2latent(style_image)
        self.content_latent = None
        null_embeds = self.encode_prompt("", device, 1, False)[0]
        self.null_embeds = null_embeds
        self.null_embeds_for_latents = torch.cat([null_embeds] * latents.shape[0])
        self.null_embeds_for_style = torch.cat(
            [null_embeds] * self.style_latent.shape[0]
        )
        self.adain = adain
        self.attn_scale = attn_scale
        self.cache = utils.DataCache()
        self.controller = controller
        utils.register_attn_control(
            self.classifier, controller=self.controller, cache=self.cache
        )
        print("Total self attention layers of Unet: ", controller.num_self_layers)
        print("Self attention layers for AD: ", controller.self_layers)

        pbar = tqdm(timesteps, desc="Sample")
        for i, t in enumerate(pbar):
            count.zero_()
            value.zero_()
            # generate views
            # Here, we iterate through different spatial crops of the latents and denoise them. These
            # denoised (latent) crops are then averaged to produce the final latent
            # for the current timestep via MultiDiffusion. Please see Sec. 4.1 in the
            # MultiDiffusion paper for more details: https://arxiv.org/abs/2302.08113
            # Batch views denoise
            for j, batch_view in enumerate(views_batch):
                vb_size = len(batch_view)
                # get the latents corresponding to the current view coordinates
                latents_for_view = torch.cat(
                    [
                        latents[:, :, h_start:h_end, w_start:w_end]
                        for h_start, h_end, w_start, w_end in batch_view
                    ]
                )
                # rematch block's scheduler status
                self.scheduler.__dict__.update(views_scheduler_status[j])

                # expand the latents if we are doing classifier free guidance
                latent_model_input = (
                    latents_for_view.repeat_interleave(2, dim=0)
                    if do_cfg
                    else latents_for_view
                )

                latent_model_input = self.scheduler.scale_model_input(
                    latent_model_input, t
                )

                # repeat prompt_embeds for batch
                prompt_embeds_input = torch.cat([prompt_embeds] * vb_size)

                # predict the noise residual
                with torch.no_grad():
                    noise_pred = self.unet(
                        latent_model_input,
                        t,
                        encoder_hidden_states=prompt_embeds_input,
                        timestep_cond=timestep_cond,
                        cross_attention_kwargs=cross_attention_kwargs,
                        added_cond_kwargs=added_cond_kwargs,
                    ).sample

                    # perform guidance
                    if do_cfg:
                        noise_pred_uncond, noise_pred_text = (
                            noise_pred[::2],
                            noise_pred[1::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_denoised_batch = self.scheduler.step(
                        noise_pred, t, latents_for_view, **extra_step_kwargs
                    ).prev_sample
                if iters > 0:
                    self.null_embeds_for_latents = torch.cat(
                        [self.null_embeds] * noise_pred.shape[0]
                    )
                    latents_denoised_batch = self.AD(
                        latents_denoised_batch, t, lr, iters, pbar
                    )
                # save views scheduler status after sample
                views_scheduler_status[j] = copy.deepcopy(self.scheduler.__dict__)

                # extract value from batch
                for latents_view_denoised, (h_start, h_end, w_start, w_end) in zip(
                    latents_denoised_batch.chunk(vb_size), batch_view
                ):

                    value[:, :, h_start:h_end, w_start:w_end] += latents_view_denoised
                    count[:, :, h_start:h_end, w_start:w_end] += 1

            # take the MultiDiffusion step. Eq. 5 in MultiDiffusion paper: https://arxiv.org/abs/2302.08113
            latents = torch.where(count > 0, value / count, value)

        images = self.latent2image(latents)
        # Offload all models
        self.maybe_free_model_hooks()
        return images

    def AD(self, latents, t, lr, iters, pbar, weight=0):
        t = max(
            t
            - self.scheduler.config.num_train_timesteps
            // self.scheduler.num_inference_steps,
            torch.tensor([0], device=self.device),
        )
        if self.adain:
            noise = torch.randn_like(self.style_latent)
            style_latent = self.scheduler.add_noise(self.style_latent, noise, t)
            latents = utils.adain(latents, style_latent)

        with torch.no_grad():
            qs_list, ks_list, vs_list, s_out_list = self.extract_feature(
                self.style_latent,
                t,
                self.null_embeds_for_style,
                add_noise=True,
            )
            if self.content_latent is not None:
                qc_list, kc_list, vc_list, c_out_list = self.extract_feature(
                    self.content_latent,
                    t,
                    self.null_embeds,
                    add_noise=True,
                )

        latents = latents.detach()
        optimizer = torch.optim.Adam([latents.requires_grad_()], lr=lr)
        optimizer = self.accelerator.prepare(optimizer)

        for j in range(iters):
            style_loss = 0
            content_loss = 0
            optimizer.zero_grad()
            q_list, k_list, v_list, self_out_list = self.extract_feature(
                latents,
                t,
                self.null_embeds_for_latents,
                add_noise=False,
            )
            style_loss = ad_loss(q_list, ks_list, vs_list, self_out_list, scale=self.attn_scale)
            if self.content_latent is not None:
                content_loss = q_loss(q_list, qc_list)
                # content_loss = qk_loss(q_list, k_list, qc_list, kc_list)
                # content_loss = qkv_loss(q_list, k_list, vc_list, c_out_list)
            loss = style_loss + content_loss * weight
            self.accelerator.backward(loss)
            optimizer.step()

            pbar.set_postfix(loss=loss.item(), time=t.item(), iter=j)
        latents = latents.detach()
        return latents

    def extract_feature(
        self,
        latent,
        t,
        embeds,
        add_noise=False,
    ):
        self.cache.clear()
        self.controller.step()
        if add_noise:
            noise = torch.randn_like(latent)
            latent_ = self.scheduler.add_noise(latent, noise, t)
        else:
            latent_ = latent
        _ = self.classifier(latent_, t, embeds)[0]
        return self.cache.get()

    def get_views_(
        self,
        panorama_height: int,
        panorama_width: int,
        window_size: int = 64,
        stride: int = 8,
    ) -> List[Tuple[int, int, int, int]]:
        panorama_height //= 8
        panorama_width //= 8

        num_blocks_height = (
            math.ceil((panorama_height - window_size) / stride) + 1
            if panorama_height > window_size
            else 1
        )
        num_blocks_width = (
            math.ceil((panorama_width - window_size) / stride) + 1
            if panorama_width > window_size
            else 1
        )

        views = []
        for i in range(int(num_blocks_height)):
            for j in range(int(num_blocks_width)):
                h_start = int(min(i * stride, panorama_height - window_size))
                w_start = int(min(j * stride, panorama_width - window_size))

                h_end = h_start + window_size
                w_end = w_start + window_size

                views.append((h_start, h_end, w_start, w_end))

        return views