From 8364ce697ddf6117fdd4f7222832d546d63880de Mon Sep 17 00:00:00 2001 From: Volpeon Date: Wed, 21 Jun 2023 13:28:49 +0200 Subject: Update --- models/attention/control.py | 106 ++++++++++++++++++++++++++++----------- models/attention/hook.py | 5 +- models/attention/structured.py | 65 ++++++++++++++---------- models/clip/embeddings.py | 29 ++++++----- models/clip/tokenizer.py | 23 +++++---- models/clip/util.py | 17 +++++-- models/convnext/discriminator.py | 11 ++-- models/sparse.py | 12 +++-- 8 files changed, 176 insertions(+), 92 deletions(-) (limited to 'models') diff --git a/models/attention/control.py b/models/attention/control.py index 248bd9f..ec378c4 100644 --- a/models/attention/control.py +++ b/models/attention/control.py @@ -23,7 +23,7 @@ class AttentionControl(abc.ABC): 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) + 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 @@ -49,12 +49,18 @@ class EmptyControl(AttentionControl): class AttentionStore(AttentionControl): @staticmethod def get_empty_store(): - return {"down_cross": [], "mid_cross": [], "up_cross": [], - "down_self": [], "mid_self": [], "up_self": []} + 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 + if attn.shape[1] <= 32**2: # avoid memory overhead self.step_store[key].append(attn) return attn @@ -68,8 +74,10 @@ class AttentionStore(AttentionControl): 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} + 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): @@ -90,7 +98,7 @@ class AttentionControlEdit(AttentionStore, abc.ABC): return x_t def replace_self_attention(self, attn_base, att_replace): - if att_replace.shape[2] <= 16 ** 2: + if att_replace.shape[2] <= 16**2: return attn_base.unsqueeze(0).expand(att_replace.shape[0], *attn_base.shape) else: return att_replace @@ -101,41 +109,62 @@ class AttentionControlEdit(AttentionStore, abc.ABC): 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]): + 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_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 - 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]): + 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], + ): super(AttentionControlEdit, self).__init__() self.batch_size = len(prompts) self.cross_replace_alpha = ptp_utils.get_time_words_attention_alpha( - prompts, num_steps, cross_replace_steps, tokenizer).to(device) + 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.num_self_replace = int(num_steps * self_replace_steps[0]), int( + num_steps * self_replace_steps[1] + ) self.local_blend = local_blend class AttentionReplace(AttentionControlEdit): def replace_cross_attention(self, attn_base, att_replace): - return torch.einsum('hpw,bwn->bhpn', attn_base, self.mapper) - - def __init__(self, prompts, num_steps: int, cross_replace_steps: float, self_replace_steps: float, - local_blend: Optional[LocalBlend] = None): - super(AttentionReplace, self).__init__(prompts, num_steps, cross_replace_steps, self_replace_steps, local_blend) + return torch.einsum("hpw,bwn->bhpn", attn_base, self.mapper) + + def __init__( + self, + prompts, + num_steps: int, + cross_replace_steps: float, + self_replace_steps: float, + local_blend: Optional[LocalBlend] = None, + ): + super(AttentionReplace, self).__init__( + prompts, num_steps, cross_replace_steps, self_replace_steps, local_blend + ) self.mapper = seq_aligner.get_replacement_mapper(prompts, tokenizer).to(device) @@ -145,9 +174,17 @@ class AttentionRefine(AttentionControlEdit): attn_replace = attn_base_replace * self.alphas + att_replace * (1 - self.alphas) return attn_replace - def __init__(self, prompts, num_steps: int, cross_replace_steps: float, self_replace_steps: float, - local_blend: Optional[LocalBlend] = None): - super(AttentionRefine, self).__init__(prompts, num_steps, cross_replace_steps, self_replace_steps, local_blend) + def __init__( + self, + prompts, + num_steps: int, + cross_replace_steps: float, + self_replace_steps: float, + local_blend: Optional[LocalBlend] = None, + ): + super(AttentionRefine, self).__init__( + prompts, num_steps, cross_replace_steps, self_replace_steps, local_blend + ) self.mapper, alphas = seq_aligner.get_refinement_mapper(prompts, tokenizer) self.mapper, alphas = self.mapper.to(device), alphas.to(device) self.alphas = alphas.reshape(alphas.shape[0], 1, 1, alphas.shape[1]) @@ -156,13 +193,24 @@ class AttentionRefine(AttentionControlEdit): class AttentionReweight(AttentionControlEdit): 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_base = self.prev_controller.replace_cross_attention( + attn_base, att_replace + ) attn_replace = attn_base[None, :, :, :] * self.equalizer[:, None, None, :] return attn_replace - 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): - super(AttentionReweight, self).__init__(prompts, num_steps, - cross_replace_steps, self_replace_steps, local_blend) + 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, + ): + super(AttentionReweight, self).__init__( + prompts, num_steps, cross_replace_steps, self_replace_steps, local_blend + ) self.equalizer = equalizer.to(device) self.prev_controller = controller diff --git a/models/attention/hook.py b/models/attention/hook.py index 903de02..6b5fb68 100644 --- a/models/attention/hook.py +++ b/models/attention/hook.py @@ -3,6 +3,7 @@ import torch try: import xformers.ops + xformers._is_functorch_available = True MEM_EFFICIENT_ATTN = True except ImportError: @@ -42,10 +43,10 @@ def register_attention_control(model, controller): return forward def register_recr(net_, count, place_in_unet): - if net_.__class__.__name__ == 'CrossAttention': + if net_.__class__.__name__ == "CrossAttention": net_.forward = ca_forward(net_, place_in_unet) return count + 1 - elif hasattr(net_, 'children'): + elif hasattr(net_, "children"): for net__ in net_.children(): count = register_recr(net__, count, place_in_unet) return count diff --git a/models/attention/structured.py b/models/attention/structured.py index 24d889f..5bbbc06 100644 --- a/models/attention/structured.py +++ b/models/attention/structured.py @@ -16,7 +16,9 @@ class StructuredAttentionControl(AttentionControl): if self.struct_attn: out = self.struct_qkv(q, context, mask) else: - context = torch.cat([context[0], context[1]['k'][0]], dim=0) # use key tensor for context + context = torch.cat( + [context[0], context[1]["k"][0]], dim=0 + ) # use key tensor for context out = self.normal_qkv(q, context, mask) else: context = default(context, x) @@ -29,11 +31,13 @@ class StructuredAttentionControl(AttentionControl): context: list of [uc, list of conditional context] """ uc_context = context[0] - context_k, context_v = context[1]['k'], context[1]['v'] + context_k, context_v = context[1]["k"], context[1]["v"] if isinstance(context_k, list) and isinstance(context_v, list): out = self.multi_qkv(q, uc_context, context_k, context_v, mask) - elif isinstance(context_k, torch.Tensor) and isinstance(context_v, torch.Tensor): + elif isinstance(context_k, torch.Tensor) and isinstance( + context_v, torch.Tensor + ): out = self.heterogeous_qkv(q, uc_context, context_k, context_v, mask) else: raise NotImplementedError @@ -50,36 +54,45 @@ class StructuredAttentionControl(AttentionControl): k_c = [self.to_k(c_k) for c_k in context_k] v_c = [self.to_v(c_v) for c_v in context_v] - q = rearrange(q, 'b n (h d) -> (b h) n d', h=h) + q = rearrange(q, "b n (h d) -> (b h) n d", h=h) - k_uc = rearrange(k_uc, 'b n (h d) -> (b h) n d', h=h) - v_uc = rearrange(v_uc, 'b n (h d) -> (b h) n d', h=h) + k_uc = rearrange(k_uc, "b n (h d) -> (b h) n d", h=h) + v_uc = rearrange(v_uc, "b n (h d) -> (b h) n d", h=h) - k_c = [rearrange(k, 'b n (h d) -> (b h) n d', h=h) for k in k_c] # NOTE: modification point - v_c = [rearrange(v, 'b n (h d) -> (b h) n d', h=h) for v in v_c] + k_c = [ + rearrange(k, "b n (h d) -> (b h) n d", h=h) for k in k_c + ] # NOTE: modification point + v_c = [rearrange(v, "b n (h d) -> (b h) n d", h=h) for v in v_c] # get composition - sim_uc = einsum('b i d, b j d -> b i j', q[:true_bs], k_uc) * self.scale - sim_c = [einsum('b i d, b j d -> b i j', q[true_bs:], k) * self.scale for k in k_c] + sim_uc = einsum("b i d, b j d -> b i j", q[:true_bs], k_uc) * self.scale + sim_c = [ + einsum("b i d, b j d -> b i j", q[true_bs:], k) * self.scale for k in k_c + ] attn_uc = sim_uc.softmax(dim=-1) attn_c = [sim.softmax(dim=-1) for sim in sim_c] # get uc output - out_uc = einsum('b i j, b j d -> b i d', attn_uc, v_uc) + out_uc = einsum("b i j, b j d -> b i d", attn_uc, v_uc) # get c output if len(v_c) == 1: out_c_collect = [] for attn in attn_c: for v in v_c: - out_c_collect.append(einsum('b i j, b j d -> b i d', attn, v)) + out_c_collect.append(einsum("b i j, b j d -> b i d", attn, v)) out_c = sum(out_c_collect) / len(out_c_collect) else: - out_c = sum([einsum('b i j, b j d -> b i d', attn, v) for attn, v in zip(attn_c, v_c)]) / len(v_c) + out_c = sum( + [ + einsum("b i j, b j d -> b i d", attn, v) + for attn, v in zip(attn_c, v_c) + ] + ) / len(v_c) out = torch.cat([out_uc, out_c], dim=0) - out = rearrange(out, '(b h) n d -> b n (h d)', h=h) + out = rearrange(out, "(b h) n d -> b n (h d)", h=h) return out @@ -88,21 +101,21 @@ class StructuredAttentionControl(AttentionControl): k = self.to_k(context) v = self.to_v(context) - q, k, v = map(lambda t: rearrange(t, 'b n (h d) -> (b h) n d', h=h), (q, k, v)) + q, k, v = map(lambda t: rearrange(t, "b n (h d) -> (b h) n d", h=h), (q, k, v)) - sim = einsum('b i d, b j d -> b i j', q, k) * self.scale + sim = einsum("b i d, b j d -> b i j", q, k) * self.scale if exists(mask): - mask = rearrange(mask, 'b ... -> b (...)') + mask = rearrange(mask, "b ... -> b (...)") max_neg_value = -torch.finfo(sim.dtype).max - mask = repeat(mask, 'b j -> (b h) () j', h=h) + mask = repeat(mask, "b j -> (b h) () j", h=h) sim.masked_fill_(~mask, max_neg_value) # attention, what we cannot get enough of attn = sim.softmax(dim=-1) - out = einsum('b i j, b j d -> b i d', attn, v) - out = rearrange(out, '(b h) n d -> b n (h d)', h=h) + out = einsum("b i j, b j d -> b i d", attn, v) + out = rearrange(out, "(b h) n d -> b n (h d)", h=h) return out @@ -111,21 +124,21 @@ class StructuredAttentionControl(AttentionControl): k = self.to_k(torch.cat([uc_context, context_k], dim=0)) v = self.to_v(torch.cat([uc_context, context_v], dim=0)) - q, k, v = map(lambda t: rearrange(t, 'b n (h d) -> (b h) n d', h=h), (q, k, v)) + q, k, v = map(lambda t: rearrange(t, "b n (h d) -> (b h) n d", h=h), (q, k, v)) - sim = einsum('b i d, b j d -> b i j', q, k) * self.scale + sim = einsum("b i d, b j d -> b i j", q, k) * self.scale if exists(mask): - mask = rearrange(mask, 'b ... -> b (...)') + mask = rearrange(mask, "b ... -> b (...)") max_neg_value = -torch.finfo(sim.dtype).max - mask = repeat(mask, 'b j -> (b h) () j', h=h) + mask = repeat(mask, "b j -> (b h) () j", h=h) sim.masked_fill_(~mask, max_neg_value) # attention, what we cannot get enough of attn = sim.softmax(dim=-1) - out = einsum('b i j, b j d -> b i d', attn, v) - out = rearrange(out, '(b h) n d -> b n (h d)', h=h) + out = einsum("b i j, b j d -> b i d", attn, v) + out = rearrange(out, "(b h) n d -> b n (h d)", h=h) return out def get_kv(self, context): diff --git a/models/clip/embeddings.py b/models/clip/embeddings.py index 7c7f2ac..8c3c6d4 100644 --- a/models/clip/embeddings.py +++ b/models/clip/embeddings.py @@ -14,7 +14,13 @@ from models.sparse import SparseEmbedding class ManagedCLIPTextEmbeddings(CLIPTextEmbeddings): - def __init__(self, config: CLIPTextConfig, embeddings: CLIPTextEmbeddings, alpha: int = 8, dropout: float = 0.0): + def __init__( + self, + config: CLIPTextConfig, + embeddings: CLIPTextEmbeddings, + alpha: int = 8, + dropout: float = 0.0, + ): super().__init__(config) self.position_embedding = embeddings.position_embedding @@ -28,7 +34,9 @@ class ManagedCLIPTextEmbeddings(CLIPTextEmbeddings): self.token_embedding.weight = embeddings.token_embedding.weight def resize(self, size: int): - self.token_embedding = self.token_embedding.new_resized(size, self.initializer_factor) + self.token_embedding = self.token_embedding.new_resized( + size, self.initializer_factor + ) def add_embed( self, @@ -46,7 +54,7 @@ class ManagedCLIPTextEmbeddings(CLIPTextEmbeddings): initializer = [initializer] if isinstance(initializer, list): - initializer = (initializer * len(token_ids))[:len(token_ids)] + initializer = (initializer * len(token_ids))[: len(token_ids)] with torch.no_grad(): initializer = self.get_embed(initializer) @@ -76,24 +84,21 @@ class ManagedCLIPTextEmbeddings(CLIPTextEmbeddings): def get_embed(self, input_ids: Union[list[int], torch.LongTensor]): if isinstance(input_ids, list): - input_ids = torch.tensor(input_ids, device=self.token_embedding.weight.device, dtype=torch.long) + input_ids = torch.tensor( + input_ids, device=self.token_embedding.weight.device, dtype=torch.long + ) return self.token_embedding(input_ids) def patch_managed_embeddings( - text_encoder: CLIPTextModel, - alpha: int = 8, - dropout: float = 0.0 + text_encoder: CLIPTextModel, alpha: int = 8, dropout: float = 0.0 ) -> ManagedCLIPTextEmbeddings: if isinstance(text_encoder.text_model.embeddings, ManagedCLIPTextEmbeddings): return text_encoder.text_model.embeddings - + text_embeddings = ManagedCLIPTextEmbeddings( - text_encoder.config, - text_encoder.text_model.embeddings, - alpha, - dropout + text_encoder.config, text_encoder.text_model.embeddings, alpha, dropout ) text_encoder.text_model.embeddings = text_embeddings return text_embeddings diff --git a/models/clip/tokenizer.py b/models/clip/tokenizer.py index 789b525..a866641 100644 --- a/models/clip/tokenizer.py +++ b/models/clip/tokenizer.py @@ -91,18 +91,21 @@ class MultiCLIPTokenizer(CLIPTokenizer): self.vector_shuffle = shuffle_none def add_multi_tokens( - self, - new_tokens: Union[str, list[str]], - num_vectors: Union[int, list[int]] = 1 + self, new_tokens: Union[str, list[str]], num_vectors: Union[int, list[int]] = 1 ) -> Union[list[int], list[list[int]]]: if isinstance(new_tokens, list): if isinstance(num_vectors, int): num_vectors = [num_vectors] * len(new_tokens) if len(num_vectors) != len(new_tokens): - raise ValueError("Expected new_tokens and num_vectors to have the same len") + raise ValueError( + "Expected new_tokens and num_vectors to have the same len" + ) - return [self.add_multi_tokens(new_token, vecs) for new_token, vecs in zip(new_tokens, num_vectors)] + return [ + self.add_multi_tokens(new_token, vecs) + for new_token, vecs in zip(new_tokens, num_vectors) + ] if isinstance(num_vectors, list): raise ValueError("Expected num_vectors to be int for single token") @@ -129,13 +132,11 @@ class MultiCLIPTokenizer(CLIPTokenizer): return [id] def expand_ids(self, ids: list[int]): - return [ - new_id - for id in ids - for new_id in self.expand_id(id) - ] + return [new_id for id in ids for new_id in self.expand_id(id)] - def expand_batched_ids(self, input_ids: Union[list[int], list[list[int]], tuple[list[int]]]): + def expand_batched_ids( + self, input_ids: Union[list[int], list[list[int]], tuple[list[int]]] + ): if isinstance(input_ids, (list, tuple)) and isinstance(input_ids[0], list): return [self.expand_ids(batch) for batch in input_ids] else: diff --git a/models/clip/util.py b/models/clip/util.py index f94fbc7..7196bb6 100644 --- a/models/clip/util.py +++ b/models/clip/util.py @@ -5,27 +5,32 @@ import torch from transformers import CLIPTokenizer, CLIPTextModel -def unify_input_ids(tokenizer: CLIPTokenizer, input_ids: list[list[int]], max_length: Optional[int] = None): +def unify_input_ids( + tokenizer: CLIPTokenizer, + input_ids: list[list[int]], + max_length: Optional[int] = None, +): if max_length is None: return tokenizer.pad( {"input_ids": input_ids}, padding=True, pad_to_multiple_of=tokenizer.model_max_length, - return_tensors="pt" + return_tensors="pt", ) else: return tokenizer.pad( {"input_ids": input_ids}, padding="max_length", max_length=max_length, - return_tensors="pt" + return_tensors="pt", ) + def get_extended_embeddings( text_encoder: CLIPTextModel, input_ids: torch.LongTensor, position_ids: Optional[torch.LongTensor] = None, - attention_mask=None + attention_mask=None, ): model_max_length = text_encoder.config.max_position_embeddings prompts = input_ids.shape[0] @@ -36,6 +41,8 @@ def get_extended_embeddings( if attention_mask is not None: attention_mask = attention_mask.view((-1, model_max_length)) - text_embeddings = text_encoder(input_ids, position_ids=position_ids, attention_mask=attention_mask)[0] + text_embeddings = text_encoder( + input_ids, position_ids=position_ids, attention_mask=attention_mask + )[0] text_embeddings = text_embeddings.view((prompts, -1, text_embeddings.shape[2])) return text_embeddings diff --git a/models/convnext/discriminator.py b/models/convnext/discriminator.py index 571b915..5798bcf 100644 --- a/models/convnext/discriminator.py +++ b/models/convnext/discriminator.py @@ -5,7 +5,7 @@ from timm.data.constants import IMAGENET_DEFAULT_MEAN, IMAGENET_DEFAULT_STD from torch.nn import functional as F -class ConvNeXtDiscriminator(): +class ConvNeXtDiscriminator: def __init__(self, model: ConvNeXt, input_size: int) -> None: self.net = model @@ -22,8 +22,13 @@ class ConvNeXtDiscriminator(): img_mean = self.img_mean.to(device=img.device, dtype=img.dtype) img_std = self.img_std.to(device=img.device, dtype=img.dtype) - img = ((img + 1.) / 2.).sub(img_mean).div(img_std) + img = ((img + 1.0) / 2.0).sub(img_mean).div(img_std) - img = F.interpolate(img, size=(self.input_size, self.input_size), mode='bicubic', align_corners=True) + img = F.interpolate( + img, + size=(self.input_size, self.input_size), + mode="bicubic", + align_corners=True, + ) pred = self.net(img) return pred diff --git a/models/sparse.py b/models/sparse.py index bd45696..e5897c9 100644 --- a/models/sparse.py +++ b/models/sparse.py @@ -15,21 +15,25 @@ class SparseEmbedding(nn.Embedding): ): nn.Embedding.__init__(self, num_embeddings, embedding_dim, **kwargs) - self.register_buffer('trainable_ids', self.weight.new_zeros(num_embeddings, dtype=torch.long) - 1) + self.register_buffer( + "trainable_ids", self.weight.new_zeros(num_embeddings, dtype=torch.long) - 1 + ) self.trainable = nn.ParameterList() self.scaling = alpha self.dropout_p = dropout self.weight.requires_grad = False - if dropout > 0.: + if dropout > 0.0: self.dropout = nn.Dropout(p=dropout) else: self.dropout = nn.Identity() self.reset_parameters() - def new_resized(self, new_num_embeddings: int, initializer_factor: Optional[float] = None): + def new_resized( + self, new_num_embeddings: int, initializer_factor: Optional[float] = None + ): n = min(self.num_embeddings, new_num_embeddings) new_emb = SparseEmbedding( @@ -38,7 +42,7 @@ class SparseEmbedding(nn.Embedding): self.scaling, self.dropout_p, device=self.weight.device, - dtype=self.weight.dtype + dtype=self.weight.dtype, ) if initializer_factor is not None: new_emb.weight.data.normal_(mean=0.0, std=initializer_factor * 0.02) -- cgit v1.2.3-70-g09d2