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.functional as F
3
4
5class ReplaceGrad(torch.autograd.Function):
6 @staticmethod
7 def forward(ctx, x_forward, x_backward):
8 ctx.shape = x_backward.shape
9 return x_forward
10
11 @staticmethod
12 def backward(ctx, grad_in):
13 return None, grad_in.sum_to_size(ctx.shape)
14
15
16class ClampWithGrad(torch.autograd.Function):
17 @staticmethod
18 def forward(ctx, input, min, max):
19 ctx.min = min
20 ctx.max = max
21 ctx.save_for_backward(input)
22 return input.clamp(min, max)
23
24 @staticmethod
25 def backward(ctx, grad_in):
26 (input,) = ctx.saved_tensors
27 return (
28 grad_in * (grad_in * (input - input.clamp(ctx.min, ctx.max)) >= 0),
29 None,
30 None,
31 )
32
33
34def vector_quantize(x, codebook):
35 d = (
36 x.pow(2).sum(dim=-1, keepdim=True)
37 + codebook.pow(2).sum(dim=1)
38 - 2 * x @ codebook.T
39 )
40 indices = d.argmin(-1)
41 x_q = F.one_hot(indices, codebook.shape[0]).to(d.dtype) @ codebook
42 return ReplaceGrad.apply(x_q, x)