Art generation with VQGAN + CLIP in Docker, with a simple web UI for anyone with a GPU. A simplified and expanded take on Kevin Costa's work.
aiai-artclipdockerdocker-composegpuhandcodedimagenetpythonpytorchtorchtorchvisionvqganvqgan-clip
1import torch
2import torch.nn as nn
3
4from core.taming.utils import Normalize
5
6
7class AttnBlock(nn.Module):
8 def __init__(self, in_channels):
9 super().__init__()
10 self.in_channels = in_channels
11
12 self.norm = Normalize(in_channels)
13 self.q = torch.nn.Conv2d(in_channels, in_channels, kernel_size=1, stride=1, padding=0)
14 self.k = torch.nn.Conv2d(in_channels, in_channels, kernel_size=1, stride=1, padding=0)
15 self.v = torch.nn.Conv2d(in_channels, in_channels, kernel_size=1, stride=1, padding=0)
16 self.proj_out = torch.nn.Conv2d(
17 in_channels, in_channels, kernel_size=1, stride=1, padding=0
18 )
19
20 def forward(self, x):
21 h_ = x
22 h_ = self.norm(h_)
23 q = self.q(h_)
24 k = self.k(h_)
25 v = self.v(h_)
26
27 # compute attention
28 b, c, h, w = q.shape
29 q = q.reshape(b, c, h * w)
30 q = q.permute(0, 2, 1) # b, hw, c
31 k = k.reshape(b, c, h * w) # b, c, hw
32 w_ = torch.bmm(q, k) # b, hw, hw w[b, i, j]=sum_c q[b, i, c]k[b, c, j]
33 w_ = w_ * (int(c)**(-0.5))
34 w_ = torch.nn.functional.softmax(w_, dim=2)
35
36 # attend to values
37 v = v.reshape(b, c, h * w)
38 w_ = w_.permute(0, 2, 1) # b, hw, hw (first hw of k, second of q)
39 h_ = torch.bmm(v, w_) # b, c, hw (hw of q) h_[b, c, j] = sum_i v[b, c, i] w_[b, i, j]
40 h_ = h_.reshape(b, c, h, w)
41
42 h_ = self.proj_out(h_)
43
44 return x + h_