utils/boxdiff.py

"""
This is an reimplementation boxdiff baseline for reference and comparison. It is not used in the Web UI and not enabled by default since the current attention guidance implementation (in `guidance`), which uses attention maps from multiple levels and attention transfer, seems to be more robust and coherent.

Credit: https://github.com/showlab/BoxDiff/blob/master/pipeline/sd_pipeline_boxdiff.py
"""

import torch
import torch.nn.functional as F
import math
import warnings
import gc
from collections.abc import Iterable
import utils
from . import guidance
from .attn import GaussianSmoothing

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


def _compute_max_attention_per_index(attention_maps: torch.Tensor,
                                     object_positions: List[List[int]],
                                     smooth_attentions: bool = False,
                                     sigma: float = 0.5,
                                     kernel_size: int = 3,
                                     normalize_eot: bool = False,
                                     bboxes: List[List[int]] = None,
                                     P: float = 0.2,
                                     L: int = 1,
                                     ) -> List[torch.Tensor]:
    """ Computes the maximum attention value for each of the tokens we wish to alter. """
    last_idx = -1
    assert not normalize_eot, "normalize_eot is unimplemented"

    attention_for_text = attention_maps[:, :, 1:last_idx]
    attention_for_text *= 100
    attention_for_text = F.softmax(attention_for_text, dim=-1)

    # Extract the maximum values
    max_indices_list_fg = []
    max_indices_list_bg = []
    dist_x = []
    dist_y = []

    for obj_idx, text_positions_per_obj in enumerate(object_positions):
        for text_position_per_obj in text_positions_per_obj:
            # Shift indices since we removed the first token
            image = attention_for_text[:, :, text_position_per_obj - 1]
            H, W = image.shape

            obj_mask = torch.zeros_like(image)
            corner_mask_x = torch.zeros(
                (W,), device=obj_mask.device, dtype=obj_mask.dtype)
            corner_mask_y = torch.zeros(
                (H,), device=obj_mask.device, dtype=obj_mask.dtype)

            obj_boxes = bboxes[obj_idx]

            # We support two level (one box per phrase) and three level (multiple boxes per phrase)
            if not isinstance(obj_boxes[0], Iterable):
                obj_boxes = [obj_boxes]

            for obj_box in obj_boxes:
                x_min, y_min, x_max, y_max = utils.scale_proportion(
                    obj_box, H=H, W=W)
                obj_mask[y_min: y_max, x_min: x_max] = 1

                corner_mask_x[max(x_min - L, 0): min(x_min + L + 1, W)] = 1.
                corner_mask_x[max(x_max - L, 0): min(x_max + L + 1, W)] = 1.
                corner_mask_y[max(y_min - L, 0): min(y_min + L + 1, H)] = 1.
                corner_mask_y[max(y_max - L, 0): min(y_max + L + 1, H)] = 1.

            bg_mask = 1 - obj_mask

            if smooth_attentions:
                smoothing = GaussianSmoothing(
                    channels=1, kernel_size=kernel_size, sigma=sigma, dim=2).cuda()
                input = F.pad(image.unsqueeze(0).unsqueeze(0),
                              (1, 1, 1, 1), mode='reflect')
                image = smoothing(input).squeeze(0).squeeze(0)

            # Inner-Box constraint
            k = (obj_mask.sum() * P).long()
            max_indices_list_fg.append(
                (image * obj_mask).reshape(-1).topk(k)[0].mean())

            # Outer-Box constraint
            k = (bg_mask.sum() * P).long()
            max_indices_list_bg.append(
                (image * bg_mask).reshape(-1).topk(k)[0].mean())

            # Corner Constraint
            gt_proj_x = torch.max(obj_mask, dim=0).values
            gt_proj_y = torch.max(obj_mask, dim=1).values

            # create gt according to the number L
            dist_x.append((F.l1_loss(image.max(dim=0)[
                          0], gt_proj_x, reduction='none') * corner_mask_x).mean())
            dist_y.append((F.l1_loss(image.max(dim=1)[
                          0], gt_proj_y, reduction='none') * corner_mask_y).mean())

    return max_indices_list_fg, max_indices_list_bg, dist_x, dist_y


def _compute_loss(max_attention_per_index_fg: List[torch.Tensor], max_attention_per_index_bg: List[torch.Tensor],
                  dist_x: List[torch.Tensor], dist_y: List[torch.Tensor], return_losses: bool = False) -> torch.Tensor:
    """ Computes the attend-and-excite loss using the maximum attention value for each token. """
    losses_fg = [max(0, 1. - curr_max)
                 for curr_max in max_attention_per_index_fg]
    losses_bg = [max(0, curr_max) for curr_max in max_attention_per_index_bg]
    loss = sum(losses_fg) + sum(losses_bg) + sum(dist_x) + sum(dist_y)

    # print(f"{losses_fg}, {losses_bg}, {dist_x}, {dist_y}, {loss}")

    if return_losses:
        return max(losses_fg), losses_fg
    else:
        return max(losses_fg), loss


def compute_ca_loss_boxdiff(saved_attn, bboxes, object_positions, guidance_attn_keys, ref_ca_saved_attns=None, ref_ca_last_token_only=True, ref_ca_word_token_only=False, word_token_indices=None, index=None, ref_ca_loss_weight=1.0, verbose=False, **kwargs):
    """
    v3 is equivalent to v2 but with new dictionary format for attention maps.
    The `saved_attn` is supposed to be passed to `save_attn_to_dict` in `cross_attention_kwargs` prior to computing ths loss.
    `AttnProcessor` will put attention maps into the `save_attn_to_dict`.

    `index` is the timestep.
    `ref_ca_word_token_only`: This has precedence over `ref_ca_last_token_only` (i.e., if both are enabled, we take the token from word rather than the last token).
    `ref_ca_last_token_only`: `ref_ca_saved_attn` comes from the attention map of the last token of the phrase in single object generation, so we apply it only to the last token of the phrase in overall generation if this is set to True. If set to False, `ref_ca_saved_attn` will be applied to all the text tokens.
    """
    loss = torch.tensor(0).float().cuda()
    object_number = len(bboxes)
    if object_number == 0:
        return loss

    attn_map_list = []

    for attn_key in guidance_attn_keys:
        # We only have 1 cross attention for mid.
        attn_map_integrated = saved_attn[attn_key]
        if not attn_map_integrated.is_cuda:
            attn_map_integrated = attn_map_integrated.cuda()
        # Example dimension: [20, 64, 77]
        attn_map = attn_map_integrated.squeeze(dim=0)
        attn_map_list.append(attn_map)
    # This averages both across layers and across attention heads
    attn_map = torch.cat(attn_map_list, dim=0).mean(dim=0)
    loss = add_ca_loss_per_attn_map_to_loss_boxdiff(
        loss, attn_map, object_number, bboxes, object_positions, verbose=verbose, **kwargs)

    if ref_ca_saved_attns is not None:
        warnings.warn('Attention reference loss is enabled in boxdiff mode. The original boxdiff does not have attention reference loss.')
        
        ref_loss = torch.tensor(0).float().cuda()
        ref_loss = guidance.add_ref_ca_loss_per_attn_map_to_lossv2(
            ref_loss, saved_attn=saved_attn, object_number=object_number, bboxes=bboxes, object_positions=object_positions, guidance_attn_keys=guidance_attn_keys,
            ref_ca_saved_attns=ref_ca_saved_attns, ref_ca_last_token_only=ref_ca_last_token_only, ref_ca_word_token_only=ref_ca_word_token_only, word_token_indices=word_token_indices, verbose=verbose, index=index, loss_weight=ref_ca_loss_weight
        )
        print(f"loss {loss.item():.3f}, reference attention loss (weighted) {ref_loss.item():.3f}")
        loss += ref_loss

    return loss


def add_ca_loss_per_attn_map_to_loss_boxdiff(original_loss, attention_maps, object_number, bboxes, object_positions, P=0.2, L=1, smooth_attentions=True, sigma=0.5, kernel_size=3, normalize_eot=False, verbose=False):
    # NOTE: normalize_eot is enabled in SD v2.1 in boxdiff
    i, j = attention_maps.shape
    H = W = int(math.sqrt(i))

    attention_maps = attention_maps.view(H, W, j)
    # attention_maps is aggregated cross attn map across layers and steps
    # attention_maps shape: [H, W, 77]
    max_attention_per_index_fg, max_attention_per_index_bg, dist_x, dist_y = _compute_max_attention_per_index(
        attention_maps=attention_maps,
        object_positions=object_positions,
        smooth_attentions=smooth_attentions,
        sigma=sigma,
        kernel_size=kernel_size,
        normalize_eot=normalize_eot,
        bboxes=bboxes,
        P=P,
        L=L
    )

    _, loss = _compute_loss(max_attention_per_index_fg,
                            max_attention_per_index_bg, dist_x, dist_y)

    return original_loss + loss


def latent_backward_guidance_boxdiff(scheduler, unet, cond_embeddings, index, bboxes, object_positions, t, latents, loss, amp_loss_scale=10, latent_scale=20, scale_range=(1., 0.5), max_index_step=25, cross_attention_kwargs=None, ref_ca_saved_attns=None, guidance_attn_keys=None, verbose=False, **kwargs):
    """
    amp_loss_scale: this scales the loss but will de-scale before applying for latents. This is to prevent overflow/underflow with amp, not to adjust the update step size.
    latent_scale: this scales the step size for update (scale_factor in boxdiff).
    """

    if index < max_index_step:
        saved_attn = {}
        full_cross_attention_kwargs = {
            'save_attn_to_dict': saved_attn,
            'save_keys': guidance_attn_keys,
        }

        if cross_attention_kwargs is not None:
            full_cross_attention_kwargs.update(cross_attention_kwargs)

        latents.requires_grad_(True)
        latent_model_input = latents
        latent_model_input = scheduler.scale_model_input(latent_model_input, t)

        unet(latent_model_input, t, encoder_hidden_states=cond_embeddings,
             return_cross_attention_probs=False, cross_attention_kwargs=full_cross_attention_kwargs)

        # TODO: could return the attention maps for the required blocks only and not necessarily the final output
        # update latents with guidance
        loss = compute_ca_loss_boxdiff(saved_attn=saved_attn, bboxes=bboxes, object_positions=object_positions, guidance_attn_keys=guidance_attn_keys,
                                       ref_ca_saved_attns=ref_ca_saved_attns, index=index, verbose=verbose, **kwargs) * amp_loss_scale

        if torch.isnan(loss):
            print("**Loss is NaN**")

        del full_cross_attention_kwargs, saved_attn
        # call gc.collect() here may release some memory

        grad_cond = torch.autograd.grad(
            loss.requires_grad_(True), [latents])[0]

        latents.requires_grad_(False)

        if True:
            warnings.warn("Using guidance scaled with sqrt scale")
            # According to boxdiff's implementation: https://github.com/Sierkinhane/BoxDiff/blob/16ffb677a9128128e04553a0200870a526731be0/pipeline/sd_pipeline_boxdiff.py#L616
            scale = (scale_range[0] + (scale_range[1] - scale_range[0])
                     * index / (len(scheduler.timesteps) - 1)) ** (0.5)
            latents = latents - latent_scale * scale / amp_loss_scale * grad_cond
        elif hasattr(scheduler, 'sigmas'):
            warnings.warn("Using guidance scaled with sigmas")
            scale = scheduler.sigmas[index] ** 2
            latents = latents - grad_cond * scale
        elif hasattr(scheduler, 'alphas_cumprod'):
            warnings.warn("Using guidance scaled with alphas_cumprod")
            # Scaling with classifier guidance
            alpha_prod_t = scheduler.alphas_cumprod[t]
            # Classifier guidance: https://arxiv.org/pdf/2105.05233.pdf
            # DDIM: https://arxiv.org/pdf/2010.02502.pdf
            scale = (1 - alpha_prod_t) ** (0.5)
            latents = latents - latent_scale * scale / amp_loss_scale * grad_cond
        else:
            warnings.warn("No scaling in guidance is performed")
            scale = 1
            latents = latents - grad_cond

        gc.collect()
        torch.cuda.empty_cache()

        if verbose:
            print(
                f"time index {index}, loss: {loss.item() / amp_loss_scale:.3f} (de-scaled with scale {amp_loss_scale:.1f}), latent grad scale: {scale:.3f}")

    return latents, loss