Spaces:
Running
Running
# -*- coding: utf-8 -*- | |
# Copyright (c) XiMing Xing. All rights reserved. | |
# Author: XiMing Xing | |
# Description: | |
from abc import ABC, abstractmethod | |
from typing import Optional, Union, Tuple, List, Dict | |
import torch | |
import torch.nn.functional as F | |
from .ptp_utils import (get_word_inds, get_time_words_attention_alpha) | |
from .seq_aligner import (get_replacement_mapper, get_refinement_mapper) | |
class AttentionControl(ABC): | |
def __init__(self): | |
self.cur_step = 0 | |
self.num_att_layers = -1 | |
self.cur_att_layer = 0 | |
def step_callback(self, x_t): | |
return x_t | |
def between_steps(self): | |
return | |
def num_uncond_att_layers(self): | |
return 0 | |
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: | |
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 | |
class EmptyControl(AttentionControl): | |
def forward(self, attn, is_cross: bool, place_in_unet: str): | |
return attn | |
class AttentionStore(AttentionControl): | |
def __init__(self): | |
super(AttentionStore, self).__init__() | |
self.step_store = self.get_empty_store() | |
self.attention_store = {} | |
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'}" | |
if attn.shape[1] <= 32 ** 2: # avoid memory overhead | |
self.step_store[key].append(attn) | |
return attn | |
def between_steps(self): | |
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] | |
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 = {} | |
class LocalBlend: | |
def __init__(self, | |
prompts: List[str], | |
words: [List[List[str]]], | |
tokenizer, | |
device, | |
threshold=.3, | |
max_num_words=77): | |
self.max_num_words = max_num_words | |
alpha_layers = torch.zeros(len(prompts), 1, 1, 1, 1, self.max_num_words) | |
for i, (prompt, words_) in enumerate(zip(prompts, words)): | |
if type(words_) is str: | |
words_ = [words_] | |
for word in words_: | |
ind = get_word_inds(prompt, word, tokenizer) | |
alpha_layers[i, :, :, :, :, ind] = 1 | |
self.alpha_layers = alpha_layers.to(device) | |
self.threshold = threshold | |
def __call__(self, x_t, attention_store): | |
k = 1 | |
maps = attention_store["down_cross"][2:4] + attention_store["up_cross"][:3] | |
maps = [item.reshape(self.alpha_layers.shape[0], -1, 1, 16, 16, self.max_num_words) for item in maps] | |
maps = torch.cat(maps, dim=1) | |
maps = (maps * self.alpha_layers).sum(-1).mean(1) | |
mask = F.max_pool2d(maps, (k * 2 + 1, k * 2 + 1), (1, 1), padding=(k, k)) | |
mask = F.interpolate(mask, size=(x_t.shape[2:])) | |
mask = mask / mask.max(2, keepdims=True)[0].max(3, keepdims=True)[0] | |
mask = mask.gt(self.threshold) | |
mask = (mask[:1] + mask[1:]).float() | |
x_t = x_t[:1] + mask * (x_t - x_t[:1]) | |
return x_t | |
class AttentionControlEdit(AttentionStore, 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: Optional[LocalBlend], | |
tokenizer, | |
device): | |
super(AttentionControlEdit, self).__init__() | |
self.tokenizer = tokenizer | |
self.device = device | |
self.batch_size = len(prompts) | |
self.cross_replace_alpha = get_time_words_attention_alpha(prompts, num_steps, cross_replace_steps, | |
self.tokenizer).to(self.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 # define outside | |
def step_callback(self, x_t): | |
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] <= 16 ** 2: | |
return attn_base.unsqueeze(0).expand(att_replace.shape[0], *attn_base.shape) | |
else: | |
return att_replace | |
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) | |
# FIXME not replace correctly | |
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, | |
local_blend: Optional[LocalBlend] = None, | |
tokenizer=None, | |
device=None): | |
super(AttentionReplace, self).__init__(prompts, num_steps, cross_replace_steps, self_replace_steps, | |
local_blend, tokenizer, device) | |
self.mapper = get_replacement_mapper(prompts, self.tokenizer).to(self.device) | |
def replace_cross_attention(self, attn_base, att_replace): | |
return torch.einsum('hpw,bwn->bhpn', attn_base, self.mapper) | |
class AttentionRefine(AttentionControlEdit): | |
def __init__(self, | |
prompts, | |
num_steps: int, | |
cross_replace_steps: float, | |
self_replace_steps: float, | |
local_blend: Optional[LocalBlend] = None, | |
tokenizer=None, | |
device=None): | |
super(AttentionRefine, self).__init__(prompts, num_steps, cross_replace_steps, self_replace_steps, | |
local_blend, tokenizer, device) | |
self.mapper, alphas = get_refinement_mapper(prompts, self.tokenizer) | |
self.mapper, alphas = self.mapper.to(self.device), alphas.to(self.device) | |
self.alphas = alphas.reshape(alphas.shape[0], 1, 1, alphas.shape[1]) | |
def replace_cross_attention(self, attn_base, att_replace): | |
attn_base_replace = attn_base[:, :, self.mapper].permute(2, 0, 1, 3) | |
attn_replace = attn_base_replace * self.alphas + att_replace * (1 - self.alphas) | |
return attn_replace | |
class AttentionReweight(AttentionControlEdit): | |
def __init__(self, | |
prompts, | |
num_steps: int, | |
cross_replace_steps: float, | |
self_replace_steps: float, | |
equalizer, | |
local_blend: Optional[LocalBlend] = None, | |
controller: Optional[AttentionControlEdit] = None, | |
tokenizer=None, | |
device=None): | |
super(AttentionReweight, self).__init__(prompts, num_steps, cross_replace_steps, self_replace_steps, | |
local_blend, tokenizer, device) | |
self.equalizer = equalizer.to(self.device) | |
self.prev_controller = controller | |
def replace_cross_attention(self, attn_base, att_replace): | |
if self.prev_controller is not None: | |
attn_base = self.prev_controller.replace_cross_attention(attn_base, att_replace) | |
attn_replace = attn_base[None, :, :, :] * self.equalizer[:, None, None, :] | |
return attn_replace | |
def get_equalizer(tokenizer, text: str, | |
word_select: Union[int, Tuple[int, ...]], | |
values: Union[List[float], Tuple[float, ...]]): | |
if type(word_select) is int or type(word_select) is str: | |
word_select = (word_select,) | |
equalizer = torch.ones(len(values), 77) | |
values = torch.tensor(values, dtype=torch.float32) | |
for word in word_select: | |
inds = get_word_inds(text, word, tokenizer) | |
equalizer[:, inds] = values | |
return equalizer | |