Transformer block for image-like data. First, project the input (aka embedding) and reshape to b, t, d. Then apply standard transformer action. Finally, reshape to image NEW: use_linear for more efficiency instead of the 1x1 convs
| 58 | return q * scale_q.to(q.dtype) |
| 59 | |
| 60 | class SpatialTransformerSimpleV2(nn.Module): |
| 61 | """ |
| 62 | Transformer block for image-like data. |
| 63 | First, project the input (aka embedding) |
| 64 | and reshape to b, t, d. |
| 65 | Then apply standard transformer action. |
| 66 | Finally, reshape to image |
| 67 | NEW: use_linear for more efficiency instead of the 1x1 convs |
| 68 | """ |
| 69 | |
| 70 | def __init__(self, in_channels, n_heads, d_head, |
| 71 | global_cond_dim, |
| 72 | do_self_attention=True, |
| 73 | dropout=0., |
| 74 | context_dim=None, |
| 75 | ): |
| 76 | super().__init__() |
| 77 | |
| 78 | |
| 79 | self.in_channels = in_channels |
| 80 | inner_dim = n_heads * d_head |
| 81 | self.n_heads = n_heads |
| 82 | self.d_head = d_head |
| 83 | self.do_self_attention=do_self_attention |
| 84 | |
| 85 | |
| 86 | self.x_in_norm = AdaRMSNorm(in_channels, global_cond_dim) |
| 87 | |
| 88 | #x to qkv |
| 89 | if self.do_self_attention: |
| 90 | self.x_qkv_proj = apply_wd(torch.nn.Linear(in_channels, inner_dim * 3, bias=False)) |
| 91 | else: |
| 92 | self.x_q_proj = apply_wd(torch.nn.Linear(in_channels, inner_dim, bias=False)) |
| 93 | self.x_scale = nn.Parameter(torch.full([self.n_heads], 10.0)) |
| 94 | |
| 95 | self.x_pos_emb = AxialRoPE(d_head // 2, self.n_heads) |
| 96 | |
| 97 | |
| 98 | #context to kv |
| 99 | self.cond_kv_proj = apply_wd(torch.nn.Linear(context_dim, inner_dim * 2, bias=False)) |
| 100 | self.cond_scale = nn.Parameter(torch.full([self.n_heads], 10.0)) |
| 101 | self.cond_pos_emb = AxialRoPE(d_head // 2, self.n_heads) |
| 102 | |
| 103 | self.ff = FeedForwardBlock(in_channels, d_ff=int(in_channels*2), cond_features=global_cond_dim, dropout=dropout) |
| 104 | |
| 105 | self.dropout = nn.Dropout(dropout) |
| 106 | self.proj_out = apply_wd(zero_module(nn.Linear(in_channels, inner_dim))) |
| 107 | |
| 108 | |
| 109 | def forward(self, x, pos, global_cond, context=None, context_pos=None): |
| 110 | b, c, h, w = x.shape |
| 111 | x_in = x |
| 112 | x = rearrange(x, 'b c h w -> b h w c') |
| 113 | context = rearrange(context, 'b c h w -> b h w c') |
| 114 | x = self.x_in_norm(x, global_cond) |
| 115 | |
| 116 | if self.do_self_attention: |
| 117 | #x to qkv |