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, nonlinearity
5
6
7class ResnetBlock(nn.Module):
8 def __init__(self, *, in_channels, out_channels=None, conv_shortcut=False,
9 dropout, temb_channels=512):
10 super().__init__()
11 self.in_channels = in_channels
12 out_channels = in_channels if out_channels is None else out_channels
13 self.out_channels = out_channels
14 self.use_conv_shortcut = conv_shortcut
15
16 self.norm1 = Normalize(in_channels)
17 self.conv1 = torch.nn.Conv2d(in_channels, out_channels, kernel_size=3, stride=1, padding=1)
18 if temb_channels > 0:
19 self.temb_proj = torch.nn.Linear(temb_channels,
20 out_channels)
21 self.norm2 = Normalize(out_channels)
22 self.dropout = torch.nn.Dropout(dropout)
23 self.conv2 = torch.nn.Conv2d(out_channels, out_channels, kernel_size=3, stride=1, padding=1)
24 if self.in_channels != self.out_channels:
25 if self.use_conv_shortcut:
26 self.conv_shortcut = torch.nn.Conv2d(
27 in_channels, out_channels, kernel_size=3, stride=1, padding=1
28 )
29 else:
30 self.nin_shortcut = torch.nn.Conv2d(
31 in_channels, out_channels, kernel_size=1, stride=1, padding=0
32 )
33
34 def forward(self, x, temb):
35 h = x
36 h = self.norm1(h)
37 h = nonlinearity(h)
38 h = self.conv1(h)
39
40 if temb is not None:
41 h = h + self.temb_proj(nonlinearity(temb))[:, :, None, None]
42
43 h = self.norm2(h)
44 h = nonlinearity(h)
45 h = self.dropout(h)
46 h = self.conv2(h)
47
48 if self.in_channels != self.out_channels:
49 if self.use_conv_shortcut:
50 x = self.conv_shortcut(x)
51 else:
52 x = self.nin_shortcut(x)
53
54 return x + h