234 lines
8.0 KiB
Python
234 lines
8.0 KiB
Python
import cv2
|
|
import abc
|
|
import torch
|
|
import numpy as np
|
|
from typing import List, Tuple
|
|
|
|
LOW_RESOURCE = False
|
|
|
|
|
|
def convert_preview_image(images):
|
|
# 转换图像为 torch.Tensor，并调整维度顺序为 NHWC
|
|
images_tensors = []
|
|
for img in images:
|
|
# 将 PIL.Image 转换为 numpy.ndarray
|
|
img_array = np.array(img)
|
|
# 转换 numpy.ndarray 为 torch.Tensor
|
|
img_tensor = torch.from_numpy(img_array).float() / 255.
|
|
# 转换图像格式为 CHW (如果需要)
|
|
if img_tensor.ndim == 3 and img_tensor.shape[-1] == 3:
|
|
img_tensor = img_tensor.permute(2, 0, 1)
|
|
# 添加批次维度并转换为 NHWC
|
|
img_tensor = img_tensor.unsqueeze(0).permute(0, 2, 3, 1)
|
|
images_tensors.append(img_tensor)
|
|
|
|
if len(images_tensors) > 1:
|
|
output_image = torch.cat(images_tensors, dim=0)
|
|
else:
|
|
output_image = images_tensors[0]
|
|
return output_image
|
|
|
|
|
|
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 self.num_att_layers if LOW_RESOURCE else 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 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):
|
|
self.cur_step = 0
|
|
self.num_att_layers = -1
|
|
self.cur_att_layer = 0
|
|
|
|
|
|
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'}"
|
|
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 = {}
|
|
|
|
def __init__(self):
|
|
super(AttentionStore, self).__init__()
|
|
self.step_store = self.get_empty_store()
|
|
self.attention_store = {}
|
|
|
|
|
|
def register_attention_control(model, controller):
|
|
def ca_forward(self, place_in_unet):
|
|
to_out = self.to_out
|
|
if type(to_out) is torch.nn.modules.container.ModuleList:
|
|
to_out = self.to_out[0]
|
|
else:
|
|
to_out = self.to_out
|
|
|
|
def forward(hidden_states, encoder_hidden_states=None, attention_mask=None, temb=None, ):
|
|
is_cross = encoder_hidden_states is not None
|
|
|
|
residual = hidden_states
|
|
|
|
if self.spatial_norm is not None:
|
|
hidden_states = self.spatial_norm(hidden_states, temb)
|
|
|
|
input_ndim = hidden_states.ndim
|
|
|
|
if input_ndim == 4:
|
|
batch_size, channel, height, width = hidden_states.shape
|
|
hidden_states = hidden_states.view(batch_size, channel, height * width).transpose(1, 2)
|
|
|
|
batch_size, sequence_length, _ = (
|
|
hidden_states.shape if encoder_hidden_states is None else encoder_hidden_states.shape
|
|
)
|
|
attention_mask = self.prepare_attention_mask(attention_mask, sequence_length, batch_size)
|
|
|
|
if self.group_norm is not None:
|
|
hidden_states = self.group_norm(hidden_states.transpose(1, 2)).transpose(1, 2)
|
|
|
|
query = self.to_q(hidden_states)
|
|
|
|
if encoder_hidden_states is None:
|
|
encoder_hidden_states = hidden_states
|
|
elif self.norm_cross:
|
|
encoder_hidden_states = self.norm_encoder_hidden_states(encoder_hidden_states)
|
|
|
|
key = self.to_k(encoder_hidden_states)
|
|
value = self.to_v(encoder_hidden_states)
|
|
|
|
query = self.head_to_batch_dim(query)
|
|
key = self.head_to_batch_dim(key)
|
|
value = self.head_to_batch_dim(value)
|
|
|
|
attention_probs = self.get_attention_scores(query, key, attention_mask)
|
|
attention_probs = controller(attention_probs, is_cross, place_in_unet)
|
|
|
|
hidden_states = torch.bmm(attention_probs, value)
|
|
hidden_states = self.batch_to_head_dim(hidden_states)
|
|
|
|
# linear proj
|
|
hidden_states = to_out(hidden_states)
|
|
|
|
if input_ndim == 4:
|
|
hidden_states = hidden_states.transpose(-1, -2).reshape(batch_size, channel, height, width)
|
|
|
|
if self.residual_connection:
|
|
hidden_states = hidden_states + residual
|
|
|
|
hidden_states = hidden_states / self.rescale_output_factor
|
|
|
|
return hidden_states
|
|
|
|
return forward
|
|
|
|
class DummyController:
|
|
|
|
def __call__(self, *args):
|
|
return args[0]
|
|
|
|
def __init__(self):
|
|
self.num_att_layers = 0
|
|
|
|
if controller is None:
|
|
controller = DummyController()
|
|
|
|
def register_recr(net_, count, place_in_unet):
|
|
if net_.__class__.__name__ == 'Attention':
|
|
net_.forward = ca_forward(net_, place_in_unet)
|
|
return count + 1
|
|
elif hasattr(net_, 'children'):
|
|
for net__ in net_.children():
|
|
count = register_recr(net__, count, place_in_unet)
|
|
return count
|
|
|
|
cross_att_count = 0
|
|
sub_nets = model.unet.named_children()
|
|
for net in sub_nets:
|
|
if "down" in net[0]:
|
|
cross_att_count += register_recr(net[1], 0, "down")
|
|
elif "up" in net[0]:
|
|
cross_att_count += register_recr(net[1], 0, "up")
|
|
elif "mid" in net[0]:
|
|
cross_att_count += register_recr(net[1], 0, "mid")
|
|
|
|
controller.num_att_layers = cross_att_count
|
|
|
|
|
|
def aggregate_attention(attention_store: AttentionStore, prompts: str, res: int, from_where: List[str], is_cross: bool,
|
|
select: int):
|
|
out = []
|
|
attention_maps = attention_store.get_average_attention()
|
|
num_pixels = res ** 2
|
|
for location in from_where:
|
|
for item in attention_maps[f"{location}_{'cross' if is_cross else 'self'}"]:
|
|
if item.shape[1] == num_pixels:
|
|
cross_maps = item.reshape(len(prompts), -1, res, res, item.shape[-1])[select]
|
|
out.append(cross_maps)
|
|
out = torch.cat(out, dim=0)
|
|
out = out.sum(0) / out.shape[0]
|
|
return out.cpu()
|
|
|
|
|
|
def text_under_image(image: np.ndarray, text: str, text_color: Tuple[int, int, int] = (0, 0, 0)):
|
|
h, w, c = image.shape
|
|
offset = int(h * .2)
|
|
img = np.ones((h + offset, w, c), dtype=np.uint8) * 255
|
|
font = cv2.FONT_HERSHEY_SIMPLEX
|
|
# font = ImageFont.truetype("/usr/share/fonts/truetype/noto/NotoMono-Regular.ttf", font_size)
|
|
img[:h] = image
|
|
textsize = cv2.getTextSize(text, font, 1, 2)[0]
|
|
text_x, text_y = (w - textsize[0]) // 2, h + offset - textsize[1] // 2
|
|
cv2.putText(img, text, (text_x, text_y), font, 1, text_color, 2)
|
|
return img
|