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
3import torch.nn.functional as F
4
5from core.utils.gradients import ReplaceGrad
6
7
8class Prompt(nn.Module):
9 def __init__(self, embed, weight=1.0, stop=float("-inf")):
10 super().__init__()
11 self.register_buffer("embed", embed)
12 self.register_buffer("weight", torch.as_tensor(weight))
13 self.register_buffer("stop", torch.as_tensor(stop))
14
15 def forward(self, input):
16 input_normed = F.normalize(input.unsqueeze(1), dim=2)
17 embed_normed = F.normalize(self.embed.unsqueeze(0), dim=2)
18 dists = input_normed.sub(embed_normed).norm(dim=2).div(2).arcsin().pow(2).mul(2)
19 dists = dists * self.weight.sign()
20 return (
21 self.weight.abs()
22 * ReplaceGrad.apply(dists, torch.maximum(dists, self.stop)).mean()
23 )
24
25
26def parse_prompt(prompt):
27 vals = prompt.rsplit(":", 2)
28 vals = vals + ["", "1", "-inf"][len(vals) :]
29 return vals[0], float(vals[1]), float(vals[2])