|
from typing import Optional, Union, Tuple, List, Callable, Dict |
|
import torch |
|
import torch.nn.functional as nnf |
|
import numpy as np |
|
import abc |
|
import src.prompt_attention.p2p_utils as p2p_utils |
|
import src.prompt_attention.seq_aligner as seq_aligner |
|
|
|
|
|
|
|
class AttentionControl(abc.ABC): |
|
|
|
def step_callback(self, x_t): |
|
return x_t |
|
|
|
def between_steps(self): |
|
return |
|
|
|
@property |
|
def num_uncond_att_layers(self): |
|
|
|
return 0 |
|
|
|
@abc.abstractmethod |
|
def forward(self, attn, is_cross: bool, place_in_unet: str): |
|
raise NotImplementedError |
|
|
|
def __call__(self, attn, is_cross: bool, place_in_unet: str): |
|
if self.cur_att_layer >= self.num_uncond_att_layers: |
|
if self.low_resource: |
|
attn = self.forward(attn, is_cross, place_in_unet) |
|
else: |
|
h = attn.shape[0] |
|
attn[h // 2:] = self.forward(attn[h // 2:], is_cross, place_in_unet) |
|
self.cur_att_layer += 1 |
|
if self.cur_att_layer == self.num_att_layers + self.num_uncond_att_layers: |
|
self.cur_att_layer = 0 |
|
self.cur_step += 1 |
|
self.between_steps() |
|
return attn |
|
|
|
def reset(self): |
|
self.cur_step = 0 |
|
self.cur_att_layer = 0 |
|
|
|
def __init__(self, low_resource=False, width=None, height=None): |
|
self.cur_step = 0 |
|
self.num_att_layers = -1 |
|
self.cur_att_layer = 0 |
|
self.low_resource = low_resource |
|
self.width = width |
|
self.height = height |
|
|
|
class AttentionStore(AttentionControl): |
|
|
|
@staticmethod |
|
def get_empty_store(): |
|
return {"down_cross": [], "mid_cross": [], "up_cross": [], |
|
"down_self": [], "mid_self": [], "up_self": []} |
|
|
|
def forward(self, attn, is_cross: bool, place_in_unet: str): |
|
key = f"{place_in_unet}_{'cross' if is_cross else 'self'}" |
|
|
|
return attn |
|
|
|
def between_steps(self): |
|
if self.save_global_store: |
|
if len(self.attention_store) == 0: |
|
self.attention_store = self.step_store |
|
else: |
|
for key in self.attention_store: |
|
for i in range(len(self.attention_store[key])): |
|
self.attention_store[key][i] += self.step_store[key][i] |
|
else: |
|
self.attention_store = self.step_store |
|
self.step_store = self.get_empty_store() |
|
|
|
def get_average_attention(self): |
|
average_attention = {key: [item / self.cur_step for item in self.attention_store[key]] for key in |
|
self.attention_store} |
|
return average_attention |
|
|
|
def reset(self): |
|
super(AttentionStore, self).reset() |
|
self.step_store = self.get_empty_store() |
|
self.attention_store = {} |
|
|
|
def __init__(self, width, height, low_resolution=False, save_global_store=False): |
|
super(AttentionStore, self).__init__(low_resolution, width, height) |
|
self.step_store = self.get_empty_store() |
|
self.attention_store = {} |
|
self.save_global_store = save_global_store |
|
|
|
class AttentionControlEdit(AttentionStore, abc.ABC): |
|
def __init__(self, prompts, num_steps: int, |
|
cross_replace_steps: Union[float, Tuple[float, float], Dict[str, Tuple[float, float]]], |
|
self_replace_steps: Union[float, Tuple[float, float]], |
|
local_blend=None, width=None, height=None, tokenizer=None, device=None): |
|
super(AttentionControlEdit, self).__init__(width, height) |
|
self.batch_size = len(prompts) |
|
self.cross_replace_alpha = p2p_utils.get_time_words_attention_alpha(prompts, num_steps, cross_replace_steps, |
|
tokenizer).to(device) |
|
if type(self_replace_steps) is float: |
|
self_replace_steps = 0, self_replace_steps |
|
self.num_self_replace = int(num_steps * self_replace_steps[0]), int(num_steps * self_replace_steps[1]) |
|
self.local_blend = local_blend |
|
|
|
def step_callback(self, x_t): |
|
print("step_callback") |
|
if self.local_blend is not None: |
|
x_t = self.local_blend(x_t, self.attention_store) |
|
return x_t |
|
|
|
def replace_self_attention(self, attn_base, att_replace): |
|
if att_replace.shape[2] <= self.width * self.height: |
|
return attn_base.unsqueeze(0).expand(att_replace.shape[0], *attn_base.shape) |
|
else: |
|
return att_replace |
|
|
|
@abc.abstractmethod |
|
def replace_cross_attention(self, attn_base, att_replace): |
|
raise NotImplementedError |
|
|
|
def forward(self, attn, is_cross: bool, place_in_unet: str): |
|
super(AttentionControlEdit, self).forward(attn, is_cross, place_in_unet) |
|
if is_cross or (self.num_self_replace[0] <= self.cur_step < self.num_self_replace[1]): |
|
h = attn.shape[0] // (self.batch_size) |
|
attn = attn.reshape(self.batch_size, h, *attn.shape[1:]) |
|
attn_base, attn_repalce = attn[0], attn[1:] |
|
if is_cross: |
|
alpha_words = self.cross_replace_alpha[self.cur_step] |
|
attn_repalce_new = self.replace_cross_attention(attn_base, attn_repalce) * alpha_words + ( |
|
1 - alpha_words) * attn_repalce |
|
attn[1:] = attn_repalce_new |
|
else: |
|
attn[1:] = self.replace_self_attention(attn_base, attn_repalce) |
|
attn = attn.reshape(self.batch_size * h, *attn.shape[2:]) |
|
return attn |
|
|
|
class AttentionReplace(AttentionControlEdit): |
|
def __init__(self, prompts, num_steps: int, cross_replace_steps: float, self_replace_steps: float, width, height, |
|
local_blend = None, tokenizer=None, device=None, dtype=None): |
|
super(AttentionReplace, self).__init__(prompts, num_steps, cross_replace_steps, self_replace_steps, local_blend, width, height, tokenizer=tokenizer, device=device) |
|
self.mapper = seq_aligner.get_replacement_mapper(prompts, tokenizer).to(dtype=dtype, device=device) |
|
|
|
def replace_cross_attention(self, attn_base, att_replace): |
|
return torch.einsum('hpw,bwn->bhpn', attn_base, self.mapper) |
|
|
|
|
