Tony Lian
Update: add attention guidance and refactor the code
89f6983
"""
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