Add files via upload
add
This commit is contained in:
smthemex
@@ -0,0 +1,494 @@
|
||||
# modified from https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention_processor.py
|
||||
# and https://github.com/tencent-ailab/IP-Adapter/blob/main/ip_adapter/attention_processor.py
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
|
||||
|
||||
def minmax_normalize(batch_maps):
|
||||
min_val = batch_maps.min(dim=-1, keepdim=True)[0].min(dim=-2, keepdim=True)[0]
|
||||
max_val = batch_maps.max(dim=-1, keepdim=True)[0].max(dim=-2, keepdim=True)[0]
|
||||
|
||||
return (batch_maps - min_val) / (max_val - min_val + 1e-5)
|
||||
|
||||
|
||||
class AttnProcessor2_0(torch.nn.Module):
|
||||
r"""
|
||||
Processor for implementing scaled dot-product attention (enabled by default if you're using PyTorch 2.0).
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
hidden_size=None,
|
||||
cross_attention_dim=None,
|
||||
):
|
||||
super().__init__()
|
||||
if not hasattr(F, "scaled_dot_product_attention"):
|
||||
raise ImportError("AttnProcessor2_0 requires PyTorch 2.0, to use it, please upgrade PyTorch to 2.0.")
|
||||
|
||||
def __call__(
|
||||
self,
|
||||
attn,
|
||||
hidden_states,
|
||||
encoder_hidden_states=None,
|
||||
attention_mask=None,
|
||||
temb=None,
|
||||
boxes=None,
|
||||
phrase_idxes=None,
|
||||
eot_idxes=None,
|
||||
):
|
||||
residual = hidden_states
|
||||
|
||||
if attn.spatial_norm is not None:
|
||||
hidden_states = attn.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
|
||||
)
|
||||
|
||||
if attention_mask is not None:
|
||||
attention_mask = attn.prepare_attention_mask(attention_mask, sequence_length, batch_size)
|
||||
# scaled_dot_product_attention expects attention_mask shape to be
|
||||
# (batch, heads, source_length, target_length)
|
||||
attention_mask = attention_mask.view(batch_size, attn.heads, -1, attention_mask.shape[-1])
|
||||
|
||||
if attn.group_norm is not None:
|
||||
hidden_states = attn.group_norm(hidden_states.transpose(1, 2)).transpose(1, 2)
|
||||
|
||||
query = attn.to_q(hidden_states)
|
||||
|
||||
if encoder_hidden_states is None:
|
||||
encoder_hidden_states = hidden_states
|
||||
elif attn.norm_cross:
|
||||
encoder_hidden_states = attn.norm_encoder_hidden_states(encoder_hidden_states)
|
||||
|
||||
key = attn.to_k(encoder_hidden_states)
|
||||
value = attn.to_v(encoder_hidden_states)
|
||||
|
||||
inner_dim = key.shape[-1]
|
||||
head_dim = inner_dim // attn.heads
|
||||
|
||||
query = query.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
|
||||
|
||||
key = key.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
|
||||
value = value.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
|
||||
|
||||
# the output of sdp = (batch, num_heads, seq_len, head_dim)
|
||||
# TODO: add support for attn.scale when we move to Torch 2.1
|
||||
hidden_states = F.scaled_dot_product_attention(
|
||||
query, key, value, attn_mask=attention_mask, dropout_p=0.0, is_causal=False
|
||||
)
|
||||
|
||||
hidden_states = hidden_states.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim)
|
||||
hidden_states = hidden_states.to(query.dtype)
|
||||
|
||||
# linear proj
|
||||
hidden_states = attn.to_out[0](hidden_states)
|
||||
# dropout
|
||||
hidden_states = attn.to_out[1](hidden_states)
|
||||
|
||||
if input_ndim == 4:
|
||||
hidden_states = hidden_states.transpose(-1, -2).reshape(batch_size, channel, height, width)
|
||||
|
||||
if attn.residual_connection:
|
||||
hidden_states = hidden_states + residual
|
||||
|
||||
hidden_states = hidden_states / attn.rescale_output_factor
|
||||
|
||||
return hidden_states
|
||||
|
||||
|
||||
class MaskedIPAttnProcessor2_0(nn.Module):
|
||||
|
||||
def __init__(self, hidden_size, cross_attention_dim=None, scale=1.0, num_tokens=4, text_tokens=77,
|
||||
need_text_attention_map=False, need_image_attention_map=True, num_dummy_tokens=4, mask_threshold=0.5,
|
||||
use_psuedo_attention_mask=False, subject_scales=None, start_step=5):
|
||||
super().__init__()
|
||||
|
||||
if not hasattr(F, "scaled_dot_product_attention"):
|
||||
raise ImportError("AttnProcessor2_0 requires PyTorch 2.0, to use it, please upgrade PyTorch to 2.0.")
|
||||
|
||||
self.hidden_size = hidden_size
|
||||
self.cross_attention_dim = cross_attention_dim
|
||||
self.scale = scale
|
||||
self.num_tokens = num_tokens
|
||||
self.text_tokens = text_tokens
|
||||
self.num_dummy_tokens = num_dummy_tokens
|
||||
self.mask_threshold = mask_threshold
|
||||
self.subject_scales = subject_scales
|
||||
self.start_step = start_step
|
||||
|
||||
self.to_k_ip = nn.Linear(cross_attention_dim or hidden_size, hidden_size, bias=False)
|
||||
self.to_v_ip = nn.Linear(cross_attention_dim or hidden_size, hidden_size, bias=False)
|
||||
|
||||
self.need_text_attention_map = need_text_attention_map
|
||||
self.need_image_attention_map = need_image_attention_map
|
||||
|
||||
self.use_psuedo_attention_mask = use_psuedo_attention_mask
|
||||
self.attention_maps = []
|
||||
|
||||
def prepare_attention_mask_qk(self, boxes, phrase_idxes, sequence_length_q, sequence_length_k,
|
||||
batch_size, head_size, dtype, device, use_masked_text_attention=False):
|
||||
if boxes is None:
|
||||
return None, None
|
||||
|
||||
# TODO: only support square image now
|
||||
num_patches_per_row = int(sequence_length_q ** 0.5)
|
||||
box_idxes_start = torch.floor(boxes[:, :, 0:2] * num_patches_per_row)
|
||||
box_idxes_end = torch.ceil(boxes[:, :, 2:4] * num_patches_per_row)
|
||||
box_idxes = torch.cat([box_idxes_start, box_idxes_end], dim=-1)
|
||||
box_masks = []
|
||||
dummy_attention_mask = torch.ones((batch_size, sequence_length_q), dtype=dtype, device=device)
|
||||
for box_idx in box_idxes.unbind(dim=1):
|
||||
x_start_patch_idx, y_start_patch_idx, x_end_patch_idx, y_end_patch_idx = box_idx.unbind(dim=1)
|
||||
x_indices = torch.arange(num_patches_per_row).unsqueeze(0).expand(batch_size, -1).to(device)
|
||||
y_indices = torch.arange(num_patches_per_row).unsqueeze(0).expand(batch_size, -1).to(device)
|
||||
x_mask = ((x_indices >= x_start_patch_idx.unsqueeze(1)) & (x_indices < x_end_patch_idx.unsqueeze(1))).to(dtype)
|
||||
y_mask = ((y_indices >= y_start_patch_idx.unsqueeze(1)) & (y_indices < y_end_patch_idx.unsqueeze(1))).to(dtype)
|
||||
box_mask = torch.bmm(y_mask.unsqueeze(2), x_mask.unsqueeze(1)).reshape(batch_size, -1)
|
||||
box_masks.append(box_mask)
|
||||
dummy_attention_mask = torch.clamp(dummy_attention_mask - box_mask, min=0)
|
||||
|
||||
# post mask
|
||||
post_dummy_attention_mask = dummy_attention_mask.to(torch.bool)
|
||||
post_dummy_attention_mask = post_dummy_attention_mask.repeat_interleave(head_size, dim=0)
|
||||
|
||||
attention_mask_qk_image = torch.stack(box_masks, dim=-1)
|
||||
attention_mask_qk_image = attention_mask_qk_image.repeat_interleave(self.num_tokens, dim=-1)
|
||||
attention_mask_qk_image = (1 - attention_mask_qk_image.to(dtype)) * -10000.0 # mask to bias
|
||||
# use dummy image tokens to process the background
|
||||
dummy_attention_mask = dummy_attention_mask.unsqueeze(-1).repeat_interleave(self.num_dummy_tokens, dim=-1)
|
||||
dummy_attention_mask = (1 - dummy_attention_mask) * -10000.0
|
||||
attention_mask_qk_image = torch.cat([dummy_attention_mask, attention_mask_qk_image], dim=-1)
|
||||
if attention_mask_qk_image.shape[0] < batch_size*head_size:
|
||||
attention_mask_qk_image = attention_mask_qk_image.repeat_interleave(head_size, dim=0)
|
||||
|
||||
if use_masked_text_attention:
|
||||
attention_mask_qk_text = torch.ones((batch_size, sequence_length_q, sequence_length_k), dtype=dtype, device=device)
|
||||
for i in range(batch_size):
|
||||
for j in range(len(box_masks)):
|
||||
start_idx, end_idx = int(phrase_idxes[i, j, 0].item()), int(phrase_idxes[i, j, 1].item())
|
||||
if start_idx == 0 and end_idx == 0:
|
||||
continue
|
||||
attention_mask_qk_text[i, :, start_idx:end_idx] = box_masks[j][i, ...].unsqueeze(-1)
|
||||
attention_mask_qk_text = (1 - attention_mask_qk_text) * -10000.0
|
||||
if attention_mask_qk_text.shape[0] < batch_size*head_size:
|
||||
attention_mask_qk_text = attention_mask_qk_text.repeat_interleave(head_size, dim=0)
|
||||
else:
|
||||
attention_mask_qk_text = None
|
||||
|
||||
return attention_mask_qk_image, attention_mask_qk_text, post_dummy_attention_mask
|
||||
|
||||
def get_text_attention_maps(self, attention_probs, boxes, phrase_idxes, head_size):
|
||||
bsz = boxes.shape[0]
|
||||
_, num_tokens_q, num_tokens_k = attention_probs.shape
|
||||
attention_probs = attention_probs.view(bsz, head_size, num_tokens_q, num_tokens_k)
|
||||
num_ref = boxes.shape[1]
|
||||
h = w = int(num_tokens_q ** 0.5)
|
||||
batch_attention_maps = []
|
||||
for i in range(bsz):
|
||||
sample_attention_maps = []
|
||||
for j in range(num_ref):
|
||||
start_idx, end_idx = int(phrase_idxes[i, j, 0].item()), int(phrase_idxes[i, j, 1].item())
|
||||
if start_idx == 0 and end_idx == 0:
|
||||
sample_attention_maps.append(
|
||||
torch.zeros(num_tokens_q, dtype=attention_probs.dtype, device=attention_probs.device))
|
||||
else:
|
||||
attention_map = attention_probs[i, :, :,
|
||||
start_idx:end_idx] # [num_heads, num_tokens_q, num_tokens_phrase]
|
||||
attention_map = torch.mean(torch.mean(attention_map, dim=-1), dim=0) # [num_tokens_q]
|
||||
sample_attention_maps.append(attention_map)
|
||||
batch_attention_maps.append(torch.stack(sample_attention_maps))
|
||||
|
||||
self.attention_maps.append(torch.stack(batch_attention_maps).reshape(bsz, num_ref, h, w))
|
||||
|
||||
def get_psuedo_attention_mask(self, head_size):
|
||||
# text_attention_maps = self.attention_maps[-1] # [bsz, num_ref, h, w]
|
||||
if not self.use_psuedo_attention_mask or len(self.attention_maps) < self.start_step:
|
||||
return None, None
|
||||
text_attention_maps = torch.stack(self.attention_maps).mean(dim=0) # [bsz, num_ref, h, w]
|
||||
text_attention_maps = minmax_normalize(text_attention_maps)
|
||||
dtype, device = text_attention_maps.dtype, text_attention_maps.device
|
||||
bsz, num_ref, h, w = text_attention_maps.shape
|
||||
seq_len_q = h * w
|
||||
text_attention_maps = text_attention_maps.view(bsz, num_ref, -1)
|
||||
text_attention_maps = text_attention_maps.transpose(1, 2) # [bsz, h*w, num_ref]
|
||||
|
||||
# use threshold to get the mask
|
||||
psuedo_attention_mask = (text_attention_maps > self.mask_threshold).to(dtype)
|
||||
psuedo_dummy_attention_mask = torch.ones((bsz, seq_len_q), dtype=dtype, device=device)
|
||||
for i in range(num_ref):
|
||||
psuedo_box_mask = psuedo_attention_mask[..., i]
|
||||
psuedo_dummy_attention_mask = torch.clamp(psuedo_dummy_attention_mask - psuedo_box_mask, min=0)
|
||||
|
||||
# post mask
|
||||
post_psuedo_dummy_attention_mask = psuedo_dummy_attention_mask.to(torch.bool)
|
||||
post_psuedo_dummy_attention_mask = post_psuedo_dummy_attention_mask.repeat_interleave(head_size, dim=0)
|
||||
|
||||
psuedo_attention_mask = psuedo_attention_mask.repeat_interleave(self.num_tokens, dim=-1)
|
||||
psuedo_attention_mask = (1 - psuedo_attention_mask) * -10000.0 # mask to bias
|
||||
psuedo_dummy_attention_mask = psuedo_dummy_attention_mask.unsqueeze(-1).repeat_interleave(self.num_dummy_tokens, dim=-1)
|
||||
psuedo_dummy_attention_mask = (1 - psuedo_dummy_attention_mask) * -10000.0
|
||||
psuedo_attention_mask = torch.cat([psuedo_dummy_attention_mask, psuedo_attention_mask], dim=-1)
|
||||
if psuedo_attention_mask.shape[0] < bsz * head_size:
|
||||
psuedo_attention_mask = psuedo_attention_mask.repeat_interleave(head_size, dim=0)
|
||||
|
||||
return psuedo_attention_mask, post_psuedo_dummy_attention_mask
|
||||
|
||||
def __call__(
|
||||
self,
|
||||
attn,
|
||||
hidden_states,
|
||||
encoder_hidden_states=None,
|
||||
attention_mask=None,
|
||||
temb=None,
|
||||
boxes=None,
|
||||
phrase_idxes=None,
|
||||
eot_idxes=None,
|
||||
):
|
||||
residual = hidden_states
|
||||
|
||||
if attn.spatial_norm is not None:
|
||||
hidden_states = attn.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
|
||||
)
|
||||
|
||||
if attention_mask is not None:
|
||||
attention_mask = attn.prepare_attention_mask(attention_mask, sequence_length, batch_size)
|
||||
# scaled_dot_product_attention expects attention_mask shape to be
|
||||
# (batch, heads, source_length, target_length)
|
||||
attention_mask = attention_mask.view(batch_size, attn.heads, -1, attention_mask.shape[-1])
|
||||
rf_attention_mask = None
|
||||
custom_attention_masks = self.prepare_attention_mask_qk(boxes, phrase_idxes, hidden_states.shape[1],
|
||||
self.text_tokens, batch_size, attn.heads,
|
||||
hidden_states.dtype, hidden_states.device,
|
||||
use_masked_text_attention=False)
|
||||
attention_mask_qk_image, attention_mask_qk_text, dummy_attention_mask = custom_attention_masks
|
||||
if attention_mask_qk_image is not None:
|
||||
attention_mask_qk_image = attention_mask_qk_image.view(batch_size, attn.heads, -1, attention_mask_qk_image.shape[-1])
|
||||
if attention_mask_qk_text is not None:
|
||||
attention_mask_qk_text = attention_mask_qk_text.view(batch_size, attn.heads, -1, attention_mask_qk_text.shape[-1])
|
||||
|
||||
if attn.group_norm is not None:
|
||||
hidden_states = attn.group_norm(hidden_states.transpose(1, 2)).transpose(1, 2)
|
||||
|
||||
query = attn.to_q(hidden_states)
|
||||
|
||||
if encoder_hidden_states is None:
|
||||
encoder_hidden_states = hidden_states
|
||||
else:
|
||||
# get encoder_hidden_states, ip_hidden_states
|
||||
# end_pos = encoder_hidden_states.shape[1] - self.num_tokens
|
||||
end_pos = self.text_tokens
|
||||
encoder_hidden_states, ip_hidden_states = (
|
||||
encoder_hidden_states[:, :end_pos, :],
|
||||
encoder_hidden_states[:, end_pos:, :],
|
||||
)
|
||||
attention_mask, rf_attention_mask = (
|
||||
attention_mask[:, :, :, :end_pos],
|
||||
attention_mask[:, :, :, end_pos:],
|
||||
) if attention_mask is not None else (None, None)
|
||||
if attn.norm_cross:
|
||||
encoder_hidden_states = attn.norm_encoder_hidden_states(encoder_hidden_states)
|
||||
|
||||
key = attn.to_k(encoder_hidden_states)
|
||||
value = attn.to_v(encoder_hidden_states)
|
||||
|
||||
inner_dim = key.shape[-1]
|
||||
head_dim = inner_dim // attn.heads
|
||||
attention_mask = attention_mask_qk_text if attention_mask_qk_text is not None else attention_mask
|
||||
if not self.need_text_attention_map:
|
||||
# original attention 2.0
|
||||
new_query = query.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
|
||||
|
||||
key = key.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
|
||||
value = value.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
|
||||
|
||||
# the output of sdp = (batch, num_heads, seq_len, head_dim)
|
||||
hidden_states = F.scaled_dot_product_attention(
|
||||
new_query, key, value, attn_mask=attention_mask, dropout_p=0.0, is_causal=False
|
||||
)
|
||||
|
||||
hidden_states = hidden_states.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim)
|
||||
hidden_states = hidden_states.to(query.dtype)
|
||||
else:
|
||||
# we need get the attention map, so use the previous attention
|
||||
new_query = attn.head_to_batch_dim(query)
|
||||
key = attn.head_to_batch_dim(key)
|
||||
value = attn.head_to_batch_dim(value)
|
||||
|
||||
if attention_mask is not None:
|
||||
attention_mask = attention_mask.view(batch_size*attn.heads, -1, attention_mask.shape[-1])
|
||||
attention_probs = attn.get_attention_scores(new_query, key, attention_mask)
|
||||
self.get_text_attention_maps(attention_probs, boxes, phrase_idxes, attn.heads)
|
||||
hidden_states = torch.bmm(attention_probs, value)
|
||||
hidden_states = attn.batch_to_head_dim(hidden_states)
|
||||
|
||||
# get psuedo attention mask for image: better start after some timesteps
|
||||
psuedo_attention_mask, psuedo_dummy_attention_mask = self.get_psuedo_attention_mask(attn.heads)
|
||||
if psuedo_attention_mask is not None:
|
||||
psuedo_attention_mask = psuedo_attention_mask.view(batch_size, attn.heads, -1,
|
||||
psuedo_attention_mask.shape[-1])
|
||||
|
||||
ip_key = self.to_k_ip(ip_hidden_states)
|
||||
ip_value = self.to_v_ip(ip_hidden_states)
|
||||
rf_attention_mask = attention_mask_qk_image if attention_mask_qk_image is not None else rf_attention_mask
|
||||
rf_attention_mask = psuedo_attention_mask if psuedo_attention_mask is not None else rf_attention_mask
|
||||
dummy_attention_mask = psuedo_dummy_attention_mask if psuedo_dummy_attention_mask is not None else dummy_attention_mask
|
||||
if not self.need_image_attention_map:
|
||||
new_query = query.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
|
||||
ip_key = ip_key.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
|
||||
ip_value = ip_value.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
|
||||
|
||||
# the output of sdp = (batch, num_heads, seq_len, head_dim)
|
||||
ip_hidden_states = F.scaled_dot_product_attention(
|
||||
new_query, ip_key, ip_value, attn_mask=rf_attention_mask, dropout_p=0.0, is_causal=False
|
||||
)
|
||||
|
||||
ip_hidden_states = ip_hidden_states.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim)
|
||||
ip_hidden_states = ip_hidden_states.to(query.dtype)
|
||||
else:
|
||||
new_query = attn.head_to_batch_dim(query)
|
||||
ip_key = attn.head_to_batch_dim(ip_key)
|
||||
ip_value = attn.head_to_batch_dim(ip_value)
|
||||
|
||||
if rf_attention_mask is not None:
|
||||
rf_attention_mask = rf_attention_mask.view(batch_size*attn.heads, -1, rf_attention_mask.shape[-1])
|
||||
ip_attention_probs = attn.get_attention_scores(new_query, ip_key, rf_attention_mask)
|
||||
# mask attention_probs in background
|
||||
ip_attention_probs = torch.where(dummy_attention_mask.unsqueeze(-1), torch.zeros_like(ip_attention_probs), ip_attention_probs)
|
||||
if self.subject_scales is not None:
|
||||
# apply different scales to different subjects
|
||||
subject_scales = torch.tensor(self.subject_scales, dtype=ip_attention_probs.dtype, device=ip_attention_probs.device)
|
||||
subject_scales = subject_scales.unsqueeze(0).unsqueeze(0).repeat_interleave(self.num_tokens, dim=-1)
|
||||
dummy_subject_scales = torch.ones((1, 1, 1), dtype=ip_attention_probs.dtype, device=ip_attention_probs.device).repeat_interleave(self.num_dummy_tokens, dim=-1)
|
||||
subject_scales = torch.cat([dummy_subject_scales, subject_scales], dim=-1)
|
||||
ip_attention_probs = ip_attention_probs * subject_scales
|
||||
ip_hidden_states = torch.bmm(ip_attention_probs, ip_value)
|
||||
ip_hidden_states = attn.batch_to_head_dim(ip_hidden_states)
|
||||
|
||||
if self.subject_scales is None:
|
||||
hidden_states = hidden_states + self.scale * ip_hidden_states
|
||||
else:
|
||||
hidden_states = hidden_states + ip_hidden_states
|
||||
|
||||
# linear proj
|
||||
hidden_states = attn.to_out[0](hidden_states)
|
||||
# dropout
|
||||
hidden_states = attn.to_out[1](hidden_states)
|
||||
|
||||
if input_ndim == 4:
|
||||
hidden_states = hidden_states.transpose(-1, -2).reshape(batch_size, channel, height, width)
|
||||
|
||||
if attn.residual_connection:
|
||||
hidden_states = hidden_states + residual
|
||||
|
||||
hidden_states = hidden_states / attn.rescale_output_factor
|
||||
|
||||
return hidden_states
|
||||
|
||||
|
||||
class CNAttnProcessor2_0:
|
||||
r"""
|
||||
Processor for implementing scaled dot-product attention (enabled by default if you're using PyTorch 2.0).
|
||||
"""
|
||||
|
||||
def __init__(self, num_tokens=4, text_tokens=77):
|
||||
if not hasattr(F, "scaled_dot_product_attention"):
|
||||
raise ImportError("AttnProcessor2_0 requires PyTorch 2.0, to use it, please upgrade PyTorch to 2.0.")
|
||||
self.num_tokens = num_tokens
|
||||
self.text_tokens = text_tokens
|
||||
|
||||
def __call__(
|
||||
self,
|
||||
attn,
|
||||
hidden_states,
|
||||
encoder_hidden_states=None,
|
||||
attention_mask=None,
|
||||
temb=None,
|
||||
):
|
||||
residual = hidden_states
|
||||
|
||||
if attn.spatial_norm is not None:
|
||||
hidden_states = attn.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
|
||||
)
|
||||
|
||||
if attention_mask is not None:
|
||||
attention_mask = attn.prepare_attention_mask(attention_mask, sequence_length, batch_size)
|
||||
# scaled_dot_product_attention expects attention_mask shape to be
|
||||
# (batch, heads, source_length, target_length)
|
||||
attention_mask = attention_mask.view(batch_size, attn.heads, -1, attention_mask.shape[-1])
|
||||
rf_attention_mask = None
|
||||
|
||||
if attn.group_norm is not None:
|
||||
hidden_states = attn.group_norm(hidden_states.transpose(1, 2)).transpose(1, 2)
|
||||
|
||||
query = attn.to_q(hidden_states)
|
||||
|
||||
if encoder_hidden_states is None:
|
||||
encoder_hidden_states = hidden_states
|
||||
else:
|
||||
# end_pos = encoder_hidden_states.shape[1] - self.num_tokens
|
||||
end_pos = self.text_tokens
|
||||
encoder_hidden_states = encoder_hidden_states[:, :end_pos] # only use text
|
||||
attention_mask = attention_mask[:, :, :end_pos]
|
||||
if attn.norm_cross:
|
||||
encoder_hidden_states = attn.norm_encoder_hidden_states(encoder_hidden_states)
|
||||
|
||||
key = attn.to_k(encoder_hidden_states)
|
||||
value = attn.to_v(encoder_hidden_states)
|
||||
|
||||
inner_dim = key.shape[-1]
|
||||
head_dim = inner_dim // attn.heads
|
||||
|
||||
query = query.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
|
||||
|
||||
key = key.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
|
||||
value = value.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
|
||||
|
||||
# the output of sdp = (batch, num_heads, seq_len, head_dim)
|
||||
# TODO: add support for attn.scale when we move to Torch 2.1
|
||||
hidden_states = F.scaled_dot_product_attention(
|
||||
query, key, value, attn_mask=attention_mask, dropout_p=0.0, is_causal=False
|
||||
)
|
||||
|
||||
hidden_states = hidden_states.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim)
|
||||
hidden_states = hidden_states.to(query.dtype)
|
||||
|
||||
# linear proj
|
||||
hidden_states = attn.to_out[0](hidden_states)
|
||||
# dropout
|
||||
hidden_states = attn.to_out[1](hidden_states)
|
||||
|
||||
if input_ndim == 4:
|
||||
hidden_states = hidden_states.transpose(-1, -2).reshape(batch_size, channel, height, width)
|
||||
|
||||
if attn.residual_connection:
|
||||
hidden_states = hidden_states + residual
|
||||
|
||||
hidden_states = hidden_states / attn.rescale_output_factor
|
||||
|
||||
return hidden_states
|
||||
@@ -0,0 +1,238 @@
|
||||
import os
|
||||
import math
|
||||
from typing import List
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
from diffusers.pipelines.controlnet import MultiControlNetModel
|
||||
|
||||
from .attention_processor import MaskedIPAttnProcessor2_0 as IPAttnProcessor, AttnProcessor2_0 as AttnProcessor, CNAttnProcessor2_0 as CNAttnProcessor
|
||||
|
||||
|
||||
class MSAdapter(torch.nn.Module):
|
||||
def __init__(self, unet, image_proj_model, adapter_modules=None, ckpt_path=None,
|
||||
num_tokens=4, text_tokens=77, max_rn=4, num_dummy_tokens=4,
|
||||
device="cuda", controlnet=None):
|
||||
super().__init__()
|
||||
self.unet = unet
|
||||
self.image_proj_model = image_proj_model
|
||||
self.adapter_modules = adapter_modules
|
||||
self.num_tokens = num_tokens
|
||||
self.num_dummy_tokens = num_dummy_tokens
|
||||
self.text_tokens = text_tokens
|
||||
self.max_rn = max_rn
|
||||
self.device = device
|
||||
self.controlnet = controlnet
|
||||
self.cross_attention_dim = self.unet.config.cross_attention_dim
|
||||
|
||||
# set attention processor when inference
|
||||
if self.adapter_modules is None:
|
||||
self.set_ms_adapter()
|
||||
|
||||
# dummy image tokens
|
||||
self.dummy_image_tokens = nn.Parameter(torch.randn(1, self.num_dummy_tokens, self.cross_attention_dim))
|
||||
|
||||
if ckpt_path is not None:
|
||||
self.load_from_checkpoint(ckpt_path)
|
||||
|
||||
def load_from_checkpoint(self, ckpt_path: str):
|
||||
if os.path.isdir(ckpt_path):
|
||||
ckpt_path = os.path.join(ckpt_path, "ms_adapter.bin")
|
||||
|
||||
# Calculate original checksums
|
||||
orig_ip_proj_sum = torch.sum(torch.stack([torch.sum(p) for p in self.image_proj_model.parameters()]))
|
||||
orig_adapter_sum = torch.sum(torch.stack([torch.sum(p) for p in self.adapter_modules.parameters()]))
|
||||
|
||||
state_dict = torch.load(ckpt_path, map_location="cpu")
|
||||
|
||||
# Load state dict for image_proj_model and adapter_modules when using resampler
|
||||
# self.image_proj_model.load_state_dict(state_dict["image_proj"], strict=False)
|
||||
# self.adapter_modules.load_state_dict(state_dict["ms_adapter"], strict=False)
|
||||
self.image_proj_model.load_state_dict(state_dict["image_proj"], strict=True)
|
||||
self.adapter_modules.load_state_dict(state_dict["ms_adapter"], strict=True)
|
||||
self.load_state_dict({"dummy_image_tokens": state_dict["dummy_image_tokens"]}, strict=False)
|
||||
|
||||
# Calculate new checksums
|
||||
new_ip_proj_sum = torch.sum(torch.stack([torch.sum(p) for p in self.image_proj_model.parameters()]))
|
||||
new_adapter_sum = torch.sum(torch.stack([torch.sum(p) for p in self.adapter_modules.parameters()]))
|
||||
|
||||
# Verify if the weights have changed
|
||||
assert orig_ip_proj_sum != new_ip_proj_sum, "Weights of image_proj_model did not change!"
|
||||
assert orig_adapter_sum != new_adapter_sum, "Weights of adapter_modules did not change!"
|
||||
|
||||
print(f"Successfully loaded weights from checkpoint {ckpt_path}")
|
||||
|
||||
def set_ms_adapter(self, weight_dtype=torch.float16, cache_attention_maps=True):
|
||||
# set attention processor
|
||||
attn_procs = {}
|
||||
for name in self.unet.attn_processors.keys():
|
||||
cross_attention_dim = None if name.endswith("attn1.processor") else self.unet.config.cross_attention_dim
|
||||
if name.startswith("mid_block"):
|
||||
hidden_size = self.unet.config.block_out_channels[-1]
|
||||
elif name.startswith("up_blocks"):
|
||||
block_id = int(name[len("up_blocks.")])
|
||||
hidden_size = list(reversed(self.unet.config.block_out_channels))[block_id]
|
||||
elif name.startswith("down_blocks"):
|
||||
block_id = int(name[len("down_blocks.")])
|
||||
hidden_size = self.unet.config.block_out_channels[block_id]
|
||||
if cross_attention_dim is None:
|
||||
attn_procs[name] = AttnProcessor()
|
||||
else:
|
||||
attn_procs[name] = IPAttnProcessor(
|
||||
hidden_size=hidden_size,
|
||||
cross_attention_dim=cross_attention_dim,
|
||||
num_tokens=self.num_tokens,
|
||||
text_tokens=self.text_tokens,
|
||||
).to(self.device, dtype=weight_dtype)
|
||||
self.unet.set_attn_processor(attn_procs)
|
||||
self.adapter_modules = torch.nn.ModuleList(self.unet.attn_processors.values())
|
||||
if self.controlnet is not None:
|
||||
if isinstance(self.controlnet, MultiControlNetModel):
|
||||
for controlnet in self.controlnet.nets:
|
||||
controlnet.set_attn_processor(CNAttnProcessor(text_tokens=self.text_tokens, num_tokens=self.num_tokens))
|
||||
else:
|
||||
self.controlnet.set_attn_processor(CNAttnProcessor(text_tokens=self.text_tokens, num_tokens=self.num_tokens))
|
||||
|
||||
@torch.inference_mode()
|
||||
def get_image_embeds(self, processed_images, image_encoder=None, image_proj_type="linear",
|
||||
image_encoder_type="clip", weight_dtype=torch.float16):
|
||||
# get image embeds
|
||||
# processed_images: [bsz, rn, ...]
|
||||
processed_images = processed_images.view(-1, processed_images.shape[-3], processed_images.shape[-2],
|
||||
processed_images.shape[-1]) # (bsz*rn, ...)
|
||||
if image_proj_type == "resampler":
|
||||
image_embeds = image_encoder(processed_images.to(self.device, dtype=weight_dtype),
|
||||
output_hidden_states=True).hidden_states[-2] # (bsz*rn, num_tokens, embedding_dim)
|
||||
else:
|
||||
image_embeds = image_encoder(processed_images.to(self.device, dtype=weight_dtype)).image_embeds # (bsz*rn, embedding_dim)
|
||||
|
||||
return image_embeds # [bsz*rn, ...]
|
||||
|
||||
def set_scale(self, scale, subject_scales):
|
||||
for attn_processor in self.pipe.unet.attn_processors.values():
|
||||
if isinstance(attn_processor, IPAttnProcessor):
|
||||
attn_processor.scale = scale
|
||||
attn_processor.subject_scales = subject_scales
|
||||
|
||||
def enable_psuedo_attention_mask(self, mask_threshold=0.5, start_step=5):
|
||||
for attn_processor in self.pipe.unet.attn_processors.values():
|
||||
if isinstance(attn_processor, IPAttnProcessor):
|
||||
attn_processor.mask_threshold = mask_threshold
|
||||
attn_processor.start_step = start_step
|
||||
attn_processor.use_psuedo_attention_mask = True
|
||||
attn_processor.need_text_attention_map = True
|
||||
attn_processor.attention_maps = [] # clear attention maps
|
||||
|
||||
def generate(self, pipe, pil_images=None, processed_images=None, prompt=None, negative_prompt=None, scale=1.0,
|
||||
num_samples=4, seed=None, guidance_scale=7.5, num_inference_steps=50, image_processor=None,
|
||||
image_encoder=None, image_proj_type="linear", image_encoder_type="clip", weight_dtype=torch.float16,
|
||||
boxes=None, phrases=None, drop_grounding_tokens=None, phrase_idxes=None,
|
||||
eot_idxes=None, height=1024, width=1024, subject_scales=None, mask_threshold=None, start_step=5,
|
||||
**kwargs):
|
||||
# generate images (validation&inference)
|
||||
self.pipe = pipe
|
||||
self.set_scale(scale, subject_scales)
|
||||
if mask_threshold is not None:
|
||||
self.enable_psuedo_attention_mask(mask_threshold, start_step)
|
||||
|
||||
# pil_images: [[xxx, xxx, xxx], [xxx, xxx, xxx], ...]
|
||||
bsz = len(pil_images) # only support bsz=1 now
|
||||
if processed_images is None:
|
||||
# write in this way to promise it can be extended to batch in the future
|
||||
processed_images = []
|
||||
for pil_image in pil_images:
|
||||
processed_image = image_processor(images=pil_image, return_tensors="pt").pixel_values
|
||||
processed_images.append(processed_image)
|
||||
processed_images = torch.stack(processed_images, dim=0)
|
||||
|
||||
num_prompts = bsz
|
||||
if prompt is None:
|
||||
prompt = "best quality, high quality"
|
||||
if negative_prompt is None:
|
||||
negative_prompt = "monochrome, lowres, bad anatomy, worst quality, low quality" # duplicate
|
||||
if not isinstance(prompt, List):
|
||||
prompt = [prompt] * num_prompts
|
||||
if not isinstance(negative_prompt, List):
|
||||
negative_prompt = [negative_prompt] * num_prompts
|
||||
|
||||
cross_attention_kwargs = None
|
||||
grounding_kwargs = None
|
||||
if boxes is not None:
|
||||
boxes = torch.tensor(boxes).to(self.device, weight_dtype)
|
||||
if phrases is not None:
|
||||
drop_grounding_tokens = drop_grounding_tokens if drop_grounding_tokens is not None else [0]*bsz
|
||||
batch_boxes = boxes.view(bsz*boxes.shape[1], -1)
|
||||
# write in this way to promise it can be extended to batch in the future
|
||||
phrase_input_ids = []
|
||||
for phrase in phrases:
|
||||
phrase_input_id = pipe.tokenizer(phrase, max_length=pipe.tokenizer.model_max_length,
|
||||
padding="max_length", truncation=True,
|
||||
return_tensors="pt").input_ids
|
||||
phrase_input_ids.append(phrase_input_id)
|
||||
phrase_input_ids = torch.stack(phrase_input_ids)
|
||||
phrase_input_ids = phrase_input_ids.view(-1, phrase_input_ids.shape[-1])
|
||||
phrase_embeds = pipe.text_encoder(phrase_input_ids.to(self.device)).pooler_output
|
||||
grounding_kwargs = {"boxes": batch_boxes, "phrase_embeds": phrase_embeds, "drop_grounding_tokens": drop_grounding_tokens}
|
||||
else:
|
||||
grounding_kwargs = None
|
||||
boxes = torch.repeat_interleave(boxes, repeats=num_samples, dim=0)
|
||||
uncond_boxes = torch.zeros_like(boxes)
|
||||
boxes = torch.cat([uncond_boxes, boxes], dim=0)
|
||||
cross_attention_kwargs = {"boxes": boxes}
|
||||
|
||||
if phrase_idxes is not None:
|
||||
phrase_idxes = torch.tensor(phrase_idxes).to(self.device, torch.int)
|
||||
eot_idxes = torch.tensor(eot_idxes).to(self.device, torch.int)
|
||||
phrase_idxes = torch.repeat_interleave(phrase_idxes, repeats=num_samples, dim=0)
|
||||
eot_idxes = torch.repeat_interleave(eot_idxes, repeats=num_samples, dim=0)
|
||||
uncond_phrase_idxes = torch.zeros_like(phrase_idxes)
|
||||
uncond_eot_idxes = torch.zeros_like(eot_idxes)
|
||||
phrase_idxes = torch.cat([uncond_phrase_idxes, phrase_idxes], dim=0)
|
||||
eot_idxes = torch.cat([uncond_eot_idxes, eot_idxes], dim=0)
|
||||
if cross_attention_kwargs is None:
|
||||
cross_attention_kwargs = {"phrase_idxes": phrase_idxes, "eot_idxes": eot_idxes}
|
||||
else:
|
||||
cross_attention_kwargs["phrase_idxes"] = phrase_idxes
|
||||
cross_attention_kwargs["eot_idxes"] = eot_idxes
|
||||
|
||||
with torch.inference_mode():
|
||||
image_embeds = self.get_image_embeds(processed_images, image_encoder, image_proj_type=image_proj_type,
|
||||
image_encoder_type=image_encoder_type, weight_dtype=weight_dtype)
|
||||
image_prompt_embeds = self.image_proj_model(image_embeds, grounding_kwargs=grounding_kwargs)
|
||||
image_prompt_embeds = image_prompt_embeds.view(bsz, -1, image_prompt_embeds.shape[-2], image_prompt_embeds.shape[-1]) # (bsz, rn, num_tokens, cross_attention_dim)
|
||||
image_prompt_embeds = image_prompt_embeds.view(bsz, image_prompt_embeds.shape[-3] * image_prompt_embeds.shape[-2],
|
||||
image_prompt_embeds.shape[-1]) # (bsz, total_num_tokens*rn, cross_attention_dim)
|
||||
image_prompt_embeds = torch.cat([self.dummy_image_tokens, image_prompt_embeds], dim=1)
|
||||
uncond_image_prompt_embeds = torch.zeros_like(image_prompt_embeds)
|
||||
bs_embed, seq_len, _ = image_prompt_embeds.shape
|
||||
image_prompt_embeds = image_prompt_embeds.repeat(1, num_samples, 1)
|
||||
image_prompt_embeds = image_prompt_embeds.view(bs_embed * num_samples, seq_len, -1)
|
||||
uncond_image_prompt_embeds = uncond_image_prompt_embeds.repeat(1, num_samples, 1)
|
||||
uncond_image_prompt_embeds = uncond_image_prompt_embeds.view(bs_embed * num_samples, seq_len, -1)
|
||||
|
||||
prompt_embeds_, negative_prompt_embeds_, pooled_prompt_embeds, negative_pooled_prompt_embeds = pipe.encode_prompt(
|
||||
prompt,
|
||||
device=self.device,
|
||||
num_images_per_prompt=num_samples,
|
||||
do_classifier_free_guidance=True,
|
||||
negative_prompt=negative_prompt,
|
||||
)
|
||||
prompt_embeds = torch.cat([prompt_embeds_, image_prompt_embeds], dim=1)
|
||||
negative_prompt_embeds = torch.cat([negative_prompt_embeds_, uncond_image_prompt_embeds], dim=1)
|
||||
|
||||
generator = torch.Generator(self.device).manual_seed(seed) if seed is not None else None
|
||||
images = pipe(
|
||||
prompt_embeds=prompt_embeds,
|
||||
negative_prompt_embeds=negative_prompt_embeds,
|
||||
pooled_prompt_embeds=pooled_prompt_embeds,
|
||||
negative_pooled_prompt_embeds=negative_pooled_prompt_embeds,
|
||||
guidance_scale=guidance_scale,
|
||||
num_inference_steps=num_inference_steps,
|
||||
generator=generator,
|
||||
cross_attention_kwargs=cross_attention_kwargs,
|
||||
height=height,
|
||||
width=width,
|
||||
**kwargs,
|
||||
).images
|
||||
|
||||
return images
|
||||
@@ -0,0 +1,257 @@
|
||||
# modified from https://github.com/mlfoundations/open_flamingo/blob/main/open_flamingo/src/helpers.py
|
||||
# and https://github.com/lucidrains/imagen-pytorch/blob/main/imagen_pytorch/imagen_pytorch.py
|
||||
# and https://github.com/tencent-ailab/IP-Adapter/blob/main/ip_adapter/ip_adapter.py
|
||||
# and https://github.com/tencent-ailab/IP-Adapter/blob/main/ip_adapter/resampler.py
|
||||
|
||||
import math
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
from einops import rearrange
|
||||
from einops.layers.torch import Rearrange
|
||||
|
||||
|
||||
class FourierEmbedder(nn.Module):
|
||||
def __init__(self, num_freqs=64, temperature=100):
|
||||
super().__init__()
|
||||
|
||||
self.num_freqs = num_freqs
|
||||
self.temperature = temperature
|
||||
|
||||
freq_bands = temperature ** (torch.arange(num_freqs) / num_freqs)
|
||||
freq_bands = freq_bands[None, None]
|
||||
self.register_buffer("freq_bands", freq_bands, persistent=False)
|
||||
|
||||
def __call__(self, x):
|
||||
x = self.freq_bands * x.unsqueeze(-1)
|
||||
return torch.stack((x.sin(), x.cos()), dim=-1).permute(0, 2, 3, 1).reshape(x.shape[0], -1)
|
||||
|
||||
|
||||
class ImageProjModel(torch.nn.Module):
|
||||
"""Projection Model"""
|
||||
|
||||
def __init__(self, cross_attention_dim=1024, clip_embeddings_dim=1024, clip_extra_context_tokens=4):
|
||||
super().__init__()
|
||||
|
||||
self.cross_attention_dim = cross_attention_dim
|
||||
self.clip_extra_context_tokens = clip_extra_context_tokens
|
||||
self.proj = torch.nn.Linear(clip_embeddings_dim, self.clip_extra_context_tokens * cross_attention_dim)
|
||||
self.norm = torch.nn.LayerNorm(cross_attention_dim)
|
||||
|
||||
def forward(self, image_embeds):
|
||||
embeds = image_embeds
|
||||
clip_extra_context_tokens = self.proj(embeds).reshape(
|
||||
-1, self.clip_extra_context_tokens, self.cross_attention_dim
|
||||
)
|
||||
clip_extra_context_tokens = self.norm(clip_extra_context_tokens)
|
||||
return clip_extra_context_tokens
|
||||
|
||||
|
||||
# FFN
|
||||
def FeedForward(dim, mult=4):
|
||||
inner_dim = int(dim * mult)
|
||||
return nn.Sequential(
|
||||
nn.LayerNorm(dim),
|
||||
nn.Linear(dim, inner_dim, bias=False),
|
||||
nn.GELU(),
|
||||
nn.Linear(inner_dim, dim, bias=False),
|
||||
# nn.LayerNorm(dim),
|
||||
)
|
||||
|
||||
|
||||
def reshape_tensor(x, heads):
|
||||
bs, length, width = x.shape
|
||||
# (bs, length, width) --> (bs, length, n_heads, dim_per_head)
|
||||
x = x.view(bs, length, heads, -1)
|
||||
# (bs, length, n_heads, dim_per_head) --> (bs, n_heads, length, dim_per_head)
|
||||
x = x.transpose(1, 2)
|
||||
# (bs, n_heads, length, dim_per_head) --> (bs*n_heads, length, dim_per_head)
|
||||
x = x.reshape(bs, heads, length, -1)
|
||||
return x
|
||||
|
||||
|
||||
class PerceiverAttention(nn.Module):
|
||||
def __init__(self, *, dim, dim_head=64, heads=8):
|
||||
super().__init__()
|
||||
self.scale = dim_head**-0.5
|
||||
self.dim_head = dim_head
|
||||
self.heads = heads
|
||||
inner_dim = dim_head * heads
|
||||
|
||||
self.norm1 = nn.LayerNorm(dim)
|
||||
self.norm2 = nn.LayerNorm(dim)
|
||||
|
||||
self.to_q = nn.Linear(dim, inner_dim, bias=False)
|
||||
self.to_kv = nn.Linear(dim, inner_dim * 2, bias=False)
|
||||
self.to_out = nn.Linear(inner_dim, dim, bias=False)
|
||||
|
||||
def forward(self, x, latents):
|
||||
"""
|
||||
Args:
|
||||
x (torch.Tensor): image features
|
||||
shape (b, n1, D)
|
||||
latent (torch.Tensor): latent features
|
||||
shape (b, n2, D)
|
||||
"""
|
||||
x = self.norm1(x)
|
||||
latents = self.norm2(latents)
|
||||
|
||||
b, l, _ = latents.shape
|
||||
|
||||
q = self.to_q(latents)
|
||||
kv_input = torch.cat((x, latents), dim=-2)
|
||||
k, v = self.to_kv(kv_input).chunk(2, dim=-1)
|
||||
|
||||
q = reshape_tensor(q, self.heads)
|
||||
k = reshape_tensor(k, self.heads)
|
||||
v = reshape_tensor(v, self.heads)
|
||||
|
||||
# attention
|
||||
scale = 1 / math.sqrt(math.sqrt(self.dim_head))
|
||||
weight = (q * scale) @ (k * scale).transpose(-2, -1) # More stable with f16 than dividing afterwards
|
||||
weight = torch.softmax(weight.float(), dim=-1).type(weight.dtype)
|
||||
out = weight @ v
|
||||
|
||||
out = out.permute(0, 2, 1, 3).reshape(b, l, -1)
|
||||
|
||||
return self.to_out(out)
|
||||
|
||||
|
||||
class Resampler(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
dim=1024,
|
||||
depth=8,
|
||||
dim_head=64,
|
||||
heads=16,
|
||||
num_queries=8,
|
||||
embedding_dim=768,
|
||||
output_dim=1024,
|
||||
ff_mult=4,
|
||||
max_seq_len: int = 257, # CLIP tokens + CLS token
|
||||
apply_pos_emb: bool = False,
|
||||
num_latents_mean_pooled: int = 0, # number of latents derived from mean pooled representation of the sequence
|
||||
latent_init_mode: str = "random",
|
||||
phrase_embeddings_dim: int = 1024,
|
||||
fourier_freqs: int = 8,
|
||||
):
|
||||
super().__init__()
|
||||
self.num_queries = num_queries
|
||||
self.grounding_token_num = self.num_queries
|
||||
self.dim = dim
|
||||
self.pos_emb = nn.Embedding(max_seq_len, embedding_dim) if apply_pos_emb else None
|
||||
|
||||
self.latent_init_mode = latent_init_mode
|
||||
if latent_init_mode == "random":
|
||||
self.latents = nn.Parameter(torch.randn(1, self.latents_token_num, dim) / dim**0.5)
|
||||
self.fourier_embedder = None
|
||||
self.latent_proj = None
|
||||
self.latent_norm = None
|
||||
elif latent_init_mode == "grounding":
|
||||
self.latents = None
|
||||
self.grounding_latents = nn.Parameter(torch.randn(1, self.grounding_token_num, dim) / dim ** 0.5)
|
||||
self.fourier_embedder = FourierEmbedder(num_freqs=fourier_freqs)
|
||||
grounding_embedding_dim = phrase_embeddings_dim + fourier_freqs * 2 * 4 # 2: sin/cos, 4: xyxy
|
||||
self.latent_proj = torch.nn.Sequential(
|
||||
torch.nn.Linear(grounding_embedding_dim, grounding_embedding_dim * 2),
|
||||
torch.nn.GELU(),
|
||||
torch.nn.Linear(grounding_embedding_dim * 2, dim * self.grounding_token_num),
|
||||
)
|
||||
self.latent_norm = nn.LayerNorm(dim)
|
||||
else:
|
||||
raise ValueError(f"Invalid latent_init_mode: {latent_init_mode}")
|
||||
|
||||
self.proj_in = nn.Linear(embedding_dim, dim)
|
||||
self.attention_norm = nn.LayerNorm(dim)
|
||||
|
||||
self.proj_out = nn.Linear(dim, output_dim)
|
||||
self.norm_out = nn.LayerNorm(output_dim)
|
||||
|
||||
self.to_latents_from_mean_pooled_seq = (
|
||||
nn.Sequential(
|
||||
nn.LayerNorm(dim),
|
||||
nn.Linear(dim, dim * num_latents_mean_pooled),
|
||||
Rearrange("b (n d) -> b n d", n=num_latents_mean_pooled),
|
||||
)
|
||||
if num_latents_mean_pooled > 0
|
||||
else None
|
||||
)
|
||||
|
||||
self.layers = nn.ModuleList([])
|
||||
for _ in range(depth):
|
||||
self.layers.append(
|
||||
nn.ModuleList(
|
||||
[
|
||||
PerceiverAttention(dim=dim, dim_head=dim_head, heads=heads),
|
||||
FeedForward(dim=dim, mult=ff_mult),
|
||||
]
|
||||
)
|
||||
)
|
||||
|
||||
def forward(self, x, grounding_kwargs=None, shortcut=False, scale=1.0):
|
||||
if self.pos_emb is not None:
|
||||
n, device = x.shape[1], x.device
|
||||
pos_emb = self.pos_emb(torch.arange(n, device=device))
|
||||
x = x + pos_emb
|
||||
|
||||
if self.latent_init_mode == "random":
|
||||
latents = self.latents.repeat(x.size(0), 1, 1)
|
||||
elif self.latent_init_mode == "grounding":
|
||||
boxes = grounding_kwargs["boxes"]
|
||||
phrase_embeds = grounding_kwargs["phrase_embeds"]
|
||||
fourier_embeds = self.fourier_embedder(boxes)
|
||||
grounding_embeds = torch.cat((phrase_embeds, fourier_embeds), dim=-1)
|
||||
|
||||
drop_grounding_tokens = grounding_kwargs["drop_grounding_tokens"]
|
||||
num_ref = x.shape[0] // len(drop_grounding_tokens)
|
||||
drop_grounding_tokens = [item for item in drop_grounding_tokens for _ in range(num_ref)]
|
||||
|
||||
latents = self.latent_proj(grounding_embeds)
|
||||
latents = latents.view(-1, self.grounding_token_num, self.dim)
|
||||
latents = self.latent_norm(latents)
|
||||
|
||||
# drop grounding tokens to learnable latents
|
||||
drop_num = len([item for item in drop_grounding_tokens if item == 1])
|
||||
if drop_num > 0:
|
||||
latents_ = []
|
||||
learnable_latents = self.grounding_latents.repeat(drop_num, 1, 1)
|
||||
cur_idx = 0
|
||||
for latent, drop_grounding_token in zip(latents, drop_grounding_tokens):
|
||||
if drop_grounding_token == 1:
|
||||
latent = learnable_latents[cur_idx]
|
||||
cur_idx += 1
|
||||
latents_.append(latent)
|
||||
latents = torch.stack(latents_)
|
||||
else:
|
||||
raise ValueError(f"Invalid latent_init_mode: {self.latent_init_mode}")
|
||||
|
||||
x = self.proj_in(x)
|
||||
|
||||
if self.to_latents_from_mean_pooled_seq:
|
||||
meanpooled_seq = masked_mean(x, dim=1, mask=torch.ones(x.shape[:2], device=x.device, dtype=torch.bool))
|
||||
meanpooled_latents = self.to_latents_from_mean_pooled_seq(meanpooled_seq)
|
||||
latents = torch.cat((meanpooled_latents, latents), dim=-2)
|
||||
|
||||
init_latents = latents
|
||||
|
||||
for attn, ff in self.layers:
|
||||
latents = attn(x, latents) + latents
|
||||
latents = ff(latents) + latents
|
||||
|
||||
latents = self.attention_norm(latents)
|
||||
latents = self.proj_out(latents)
|
||||
if shortcut:
|
||||
latents = init_latents + latents * scale
|
||||
|
||||
return self.norm_out(latents)
|
||||
|
||||
|
||||
def masked_mean(t, *, dim, mask=None):
|
||||
if mask is None:
|
||||
return t.mean(dim=dim)
|
||||
|
||||
denom = mask.sum(dim=dim, keepdim=True)
|
||||
mask = rearrange(mask, "b n -> b n 1")
|
||||
masked_t = t.masked_fill(~mask, 0.0)
|
||||
|
||||
return masked_t.sum(dim=dim) / denom.clamp(min=1e-5)
|
||||
Reference in New Issue
Block a user