| 22 | # TODO: refactor AttnBlock, CrossAttention, CLIPAttention to share code |
| 23 | |
| 24 | class AttnBlock: |
| 25 | def __init__(self, in_channels): |
| 26 | self.norm = GroupNorm(32, in_channels) |
| 27 | self.q = Conv2d(in_channels, in_channels, 1) |
| 28 | self.k = Conv2d(in_channels, in_channels, 1) |
| 29 | self.v = Conv2d(in_channels, in_channels, 1) |
| 30 | self.proj_out = Conv2d(in_channels, in_channels, 1) |
| 31 | |
| 32 | # copied from AttnBlock in ldm repo |
| 33 | def __call__(self, x): |
| 34 | h_ = self.norm(x) |
| 35 | q,k,v = self.q(h_), self.k(h_), self.v(h_) |
| 36 | |
| 37 | # compute attention |
| 38 | b,c,h,w = q.shape |
| 39 | q = q.reshape(b,c,h*w) |
| 40 | q = q.permute(0,2,1) # b,hw,c |
| 41 | k = k.reshape(b,c,h*w) # b,c,hw |
| 42 | w_ = q @ k |
| 43 | w_ = w_ * (c**(-0.5)) |
| 44 | w_ = w_.softmax() |
| 45 | |
| 46 | # attend to values |
| 47 | v = v.reshape(b,c,h*w) |
| 48 | w_ = w_.permute(0,2,1) |
| 49 | h_ = v @ w_ |
| 50 | h_ = h_.reshape(b,c,h,w) |
| 51 | |
| 52 | return x + self.proj_out(h_) |
| 53 | |
| 54 | class ResnetBlock: |
| 55 | def __init__(self, in_channels, out_channels=None): |