AnchoredAI/llm-grounded-diffusion
0
1# Copyright 2023 The HuggingFace Team. All rights reserved.2#3# Licensed under the Apache License, Version 2.0 (the "License");4# you may not use this file except in compliance with the License.5# You may obtain a copy of the License at6#7# http://www.apache.org/licenses/LICENSE-2.08#9# Unless required by applicable law or agreed to in writing, software10# distributed under the License is distributed on an "AS IS" BASIS,11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.12# See the License for the specific language governing permissions and13# limitations under the License.14from typing import Any, Dict, Optional15 16import torch17import torch.nn.functional as F18from torch import nn19 20from diffusers.utils import maybe_allow_in_graph21from .attention_processor import Attention22from diffusers.models.embeddings import CombinedTimestepLabelEmbeddings23 24# https://github.com/gligen/diffusers/blob/23a9a0fab1b48752c7b9bcc98f6fe3b1d8fa7990/src/diffusers/models/attention.py25class GatedSelfAttentionDense(nn.Module):26 def __init__(self, query_dim, context_dim, n_heads, d_head):27 super().__init__()28 29 # we need a linear projection since we need cat visual feature and obj feature30 self.linear = nn.Linear(context_dim, query_dim)31 32 self.attn = Attention(query_dim=query_dim, heads=n_heads, dim_head=d_head)33 self.ff = FeedForward(query_dim, activation_fn="geglu")34 35 self.norm1 = nn.LayerNorm(query_dim)36 self.norm2 = nn.LayerNorm(query_dim)37 38 self.register_parameter('alpha_attn', nn.Parameter(torch.tensor(0.)))39 self.register_parameter('alpha_dense', nn.Parameter(torch.tensor(0.)))40 41 self.enabled = True42 43 def forward(self, x, objs, fuser_attn_kwargs={}):44 if not self.enabled:45 return x46 47 n_visual = x.shape[1]48 objs = self.linear(objs)49 50 x = x + self.alpha_attn.tanh() * self.attn(self.norm1(torch.cat([x, objs], dim=1)), **fuser_attn_kwargs)[:, :n_visual, :]51 x = x + self.alpha_dense.tanh() * self.ff(self.norm2(x)) 52 53 return x54 55@maybe_allow_in_graph56class BasicTransformerBlock(nn.Module):57 r"""58 A basic Transformer block.59 60 Parameters:61 dim (`int`): The number of channels in the input and output.62 num_attention_heads (`int`): The number of heads to use for multi-head attention.63 attention_head_dim (`int`): The number of channels in each head.64 dropout (`float`, *optional*, defaults to 0.0): The dropout probability to use.65 cross_attention_dim (`int`, *optional*): The size of the encoder_hidden_states vector for cross attention.66 only_cross_attention (`bool`, *optional*):67 Whether to use only cross-attention layers. In this case two cross attention layers are used.68 double_self_attention (`bool`, *optional*):69 Whether to use two self-attention layers. In this case no cross attention layers are used.70 activation_fn (`str`, *optional*, defaults to `"geglu"`): Activation function to be used in feed-forward.71 num_embeds_ada_norm (:72 obj: `int`, *optional*): The number of diffusion steps used during training. See `Transformer2DModel`.73 attention_bias (:74 obj: `bool`, *optional*, defaults to `False`): Configure if the attentions should contain a bias parameter.75 """76 77 def __init__(78 self,79 dim: int,80 num_attention_heads: int,81 attention_head_dim: int,82 dropout=0.0,83 cross_attention_dim: Optional[int] = None,84 activation_fn: str = "geglu",85 num_embeds_ada_norm: Optional[int] = None,86 attention_bias: bool = False,87 only_cross_attention: bool = False,88 double_self_attention: bool = False,89 upcast_attention: bool = False,90 norm_elementwise_affine: bool = True,91 norm_type: str = "layer_norm",92 final_dropout: bool = False,93 use_gated_attention: bool = False,94 ):95 super().__init__()96 self.only_cross_attention = only_cross_attention97 98 self.use_ada_layer_norm_zero = (num_embeds_ada_norm is not None) and norm_type == "ada_norm_zero"99 self.use_ada_layer_norm = (num_embeds_ada_norm is not None) and norm_type == "ada_norm"100 101 if norm_type in ("ada_norm", "ada_norm_zero") and num_embeds_ada_norm is None:102 raise ValueError(103 f"`norm_type` is set to {norm_type}, but `num_embeds_ada_norm` is not defined. Please make sure to"104 f" define `num_embeds_ada_norm` if setting `norm_type` to {norm_type}."105 )106 107 # Define 3 blocks. Each block has its own normalization layer.108 # 1. Self-Attn109 if self.use_ada_layer_norm:110 self.norm1 = AdaLayerNorm(dim, num_embeds_ada_norm)111 elif self.use_ada_layer_norm_zero:112 self.norm1 = AdaLayerNormZero(dim, num_embeds_ada_norm)113 else:114 self.norm1 = nn.LayerNorm(dim, elementwise_affine=norm_elementwise_affine)115 self.attn1 = Attention(116 query_dim=dim,117 heads=num_attention_heads,118 dim_head=attention_head_dim,119 dropout=dropout,120 bias=attention_bias,121 cross_attention_dim=cross_attention_dim if only_cross_attention else None,122 upcast_attention=upcast_attention,123 )124 125 # 2. Cross-Attn126 if cross_attention_dim is not None or double_self_attention:127 # We currently only use AdaLayerNormZero for self attention where there will only be one attention block.128 # I.e. the number of returned modulation chunks from AdaLayerZero would not make sense if returned during129 # the second cross attention block.130 self.norm2 = (131 AdaLayerNorm(dim, num_embeds_ada_norm)132 if self.use_ada_layer_norm133 else nn.LayerNorm(dim, elementwise_affine=norm_elementwise_affine)134 )135 self.attn2 = Attention(136 query_dim=dim,137 cross_attention_dim=cross_attention_dim if not double_self_attention else None,138 heads=num_attention_heads,139 dim_head=attention_head_dim,140 dropout=dropout,141 bias=attention_bias,142 upcast_attention=upcast_attention,143 ) # is self-attn if encoder_hidden_states is none144 else:145 self.norm2 = None146 self.attn2 = None147 148 # 3. Feed-forward149 self.norm3 = nn.LayerNorm(dim, elementwise_affine=norm_elementwise_affine)150 self.ff = FeedForward(dim, dropout=dropout, activation_fn=activation_fn, final_dropout=final_dropout)151 152 # 4. Fuser153 if use_gated_attention:154 self.fuser = GatedSelfAttentionDense(dim, cross_attention_dim, num_attention_heads, attention_head_dim) 155 156 def forward(157 self,158 hidden_states: torch.FloatTensor,159 attention_mask: Optional[torch.FloatTensor] = None,160 encoder_hidden_states: Optional[torch.FloatTensor] = None,161 encoder_attention_mask: Optional[torch.FloatTensor] = None,162 timestep: Optional[torch.LongTensor] = None,163 cross_attention_kwargs: Dict[str, Any] = None,164 class_labels: Optional[torch.LongTensor] = None,165 return_cross_attention_probs: bool = None,166 ):167 # Notice that normalization is always applied before the real computation in the following blocks.168 169 # 0. Prepare GLIGEN inputs170 if 'gligen' in cross_attention_kwargs:171 cross_attention_kwargs = cross_attention_kwargs.copy() if cross_attention_kwargs is not None else {}172 gligen_kwargs = cross_attention_kwargs.pop('gligen', None)173 else:174 cross_attention_kwargs = cross_attention_kwargs if cross_attention_kwargs is not None else {}175 gligen_kwargs = None176 177 # 1. Self-Attention178 if self.use_ada_layer_norm:179 norm_hidden_states = self.norm1(hidden_states, timestep)180 elif self.use_ada_layer_norm_zero:181 norm_hidden_states, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.norm1(182 hidden_states, timestep, class_labels, hidden_dtype=hidden_states.dtype183 )184 else:185 norm_hidden_states = self.norm1(hidden_states)186 187 attn_output = self.attn1(188 norm_hidden_states,189 encoder_hidden_states=encoder_hidden_states if self.only_cross_attention else None,190 attention_mask=attention_mask,191 **cross_attention_kwargs,192 )193 if self.use_ada_layer_norm_zero:194 attn_output = gate_msa.unsqueeze(1) * attn_output195 hidden_states = attn_output + hidden_states196 197 # 1.5 GLIGEN Control198 if gligen_kwargs is not None:199 # print(gligen_kwargs)200 hidden_states = self.fuser(hidden_states, gligen_kwargs['objs'], fuser_attn_kwargs=gligen_kwargs.get("fuser_attn_kwargs", {}))201 # 1.5 ends202 203 # 2. Cross-Attention204 if self.attn2 is not None:205 norm_hidden_states = (206 self.norm2(hidden_states, timestep) if self.use_ada_layer_norm else self.norm2(hidden_states)207 )208 209 attn_output = self.attn2(210 norm_hidden_states,211 encoder_hidden_states=encoder_hidden_states,212 attention_mask=encoder_attention_mask,213 return_attntion_probs=return_cross_attention_probs,214 **cross_attention_kwargs,215 )216 217 if return_cross_attention_probs:218 attn_output, cross_attention_probs = attn_output219 220 hidden_states = attn_output + hidden_states221 222 # 3. Feed-forward223 norm_hidden_states = self.norm3(hidden_states)224 225 if self.use_ada_layer_norm_zero:226 norm_hidden_states = norm_hidden_states * (1 + scale_mlp[:, None]) + shift_mlp[:, None]227 228 ff_output = self.ff(norm_hidden_states)229 230 if self.use_ada_layer_norm_zero:231 ff_output = gate_mlp.unsqueeze(1) * ff_output232 233 hidden_states = ff_output + hidden_states234 235 if return_cross_attention_probs and self.attn2 is not None:236 return hidden_states, cross_attention_probs237 return hidden_states238 239 240class FeedForward(nn.Module):241 r"""242 A feed-forward layer.243 244 Parameters:245 dim (`int`): The number of channels in the input.246 dim_out (`int`, *optional*): The number of channels in the output. If not given, defaults to `dim`.247 mult (`int`, *optional*, defaults to 4): The multiplier to use for the hidden dimension.248 dropout (`float`, *optional*, defaults to 0.0): The dropout probability to use.249 activation_fn (`str`, *optional*, defaults to `"geglu"`): Activation function to be used in feed-forward.250 final_dropout (`bool` *optional*, defaults to False): Apply a final dropout.251 """252 253 def __init__(254 self,255 dim: int,256 dim_out: Optional[int] = None,257 mult: int = 4,258 dropout: float = 0.0,259 activation_fn: str = "geglu",260 final_dropout: bool = False,261 ):262 super().__init__()263 inner_dim = int(dim * mult)264 dim_out = dim_out if dim_out is not None else dim265 266 if activation_fn == "gelu":267 act_fn = GELU(dim, inner_dim)268 if activation_fn == "gelu-approximate":269 act_fn = GELU(dim, inner_dim, approximate="tanh")270 elif activation_fn == "geglu":271 act_fn = GEGLU(dim, inner_dim)272 elif activation_fn == "geglu-approximate":273 act_fn = ApproximateGELU(dim, inner_dim)274 275 self.net = nn.ModuleList([])276 # project in277 self.net.append(act_fn)278 # project dropout279 self.net.append(nn.Dropout(dropout))280 # project out281 self.net.append(nn.Linear(inner_dim, dim_out))282 # FF as used in Vision Transformer, MLP-Mixer, etc. have a final dropout283 if final_dropout:284 self.net.append(nn.Dropout(dropout))285 286 def forward(self, hidden_states):287 for module in self.net:288 hidden_states = module(hidden_states)289 return hidden_states290 291 292class GELU(nn.Module):293 r"""294 GELU activation function with tanh approximation support with `approximate="tanh"`.295 """296 297 def __init__(self, dim_in: int, dim_out: int, approximate: str = "none"):298 super().__init__()299 self.proj = nn.Linear(dim_in, dim_out)300 self.approximate = approximate301 302 def gelu(self, gate):303 if gate.device.type != "mps":304 return F.gelu(gate, approximate=self.approximate)305 # mps: gelu is not implemented for float16306 return F.gelu(gate.to(dtype=torch.float32), approximate=self.approximate).to(dtype=gate.dtype)307 308 def forward(self, hidden_states):309 hidden_states = self.proj(hidden_states)310 hidden_states = self.gelu(hidden_states)311 return hidden_states312 313 314class GEGLU(nn.Module):315 r"""316 A variant of the gated linear unit activation function from https://arxiv.org/abs/2002.05202.317 318 Parameters:319 dim_in (`int`): The number of channels in the input.320 dim_out (`int`): The number of channels in the output.321 """322 323 def __init__(self, dim_in: int, dim_out: int):324 super().__init__()325 self.proj = nn.Linear(dim_in, dim_out * 2)326 327 def gelu(self, gate):328 if gate.device.type != "mps":329 return F.gelu(gate)330 # mps: gelu is not implemented for float16331 return F.gelu(gate.to(dtype=torch.float32)).to(dtype=gate.dtype)332 333 def forward(self, hidden_states):334 hidden_states, gate = self.proj(hidden_states).chunk(2, dim=-1)335 return hidden_states * self.gelu(gate)336 337 338class ApproximateGELU(nn.Module):339 """340 The approximate form of Gaussian Error Linear Unit (GELU)341 342 For more details, see section 2: https://arxiv.org/abs/1606.08415343 """344 345 def __init__(self, dim_in: int, dim_out: int):346 super().__init__()347 self.proj = nn.Linear(dim_in, dim_out)348 349 def forward(self, x):350 x = self.proj(x)351 return x * torch.sigmoid(1.702 * x)352 353 354class AdaLayerNorm(nn.Module):355 """356 Norm layer modified to incorporate timestep embeddings.357 """358 359 def __init__(self, embedding_dim, num_embeddings):360 super().__init__()361 self.emb = nn.Embedding(num_embeddings, embedding_dim)362 self.silu = nn.SiLU()363 self.linear = nn.Linear(embedding_dim, embedding_dim * 2)364 self.norm = nn.LayerNorm(embedding_dim, elementwise_affine=False)365 366 def forward(self, x, timestep):367 emb = self.linear(self.silu(self.emb(timestep)))368 scale, shift = torch.chunk(emb, 2)369 x = self.norm(x) * (1 + scale) + shift370 return x371 372 373class AdaLayerNormZero(nn.Module):374 """375 Norm layer adaptive layer norm zero (adaLN-Zero).376 """377 378 def __init__(self, embedding_dim, num_embeddings):379 super().__init__()380 381 self.emb = CombinedTimestepLabelEmbeddings(num_embeddings, embedding_dim)382 383 self.silu = nn.SiLU()384 self.linear = nn.Linear(embedding_dim, 6 * embedding_dim, bias=True)385 self.norm = nn.LayerNorm(embedding_dim, elementwise_affine=False, eps=1e-6)386 387 def forward(self, x, timestep, class_labels, hidden_dtype=None):388 emb = self.linear(self.silu(self.emb(timestep, class_labels, hidden_dtype=hidden_dtype)))389 shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = emb.chunk(6, dim=1)390 x = self.norm(x) * (1 + scale_msa[:, None]) + shift_msa[:, None]391 return x, gate_msa, shift_mlp, scale_mlp, gate_mlp392 393