Pixel art generation using discrete diffusion

A developer known as litlig released a notebook demonstrating pixel art generation using discrete diffusion, a method that treats pixel art as categorical data and progressively de-masks pixels to generate 16x16 front-facing character images. The approach uses a multi-head transformer trained on a Kaggle dataset of 5 categories, with a denoise schedule alpha_t governing the jump from masked to final colors, and the model predicts the most likely ending state at each step. The code is available on GitHub and builds on prior work on language generation with discrete diffusion.

The code below is trimmed to the essentials; the full version lives in the original notebook: github.com/litlig/notebooks/pixel art discrete diffusion.ipynbFollow-up to Language generation with discrete diffusion . Diffusion models are widely used to generate images. Pixel art, with limited data points and a categorical palette, can be a good fit for discrete diffusion. The dataset https://www.kaggle.com/datasets/litlig/palette-pixel-art contains 16x16 images in 5 categories. What each category means is not explicitly stated, just inferring from the dataset, the first and the last seems to be front facing and side facing of characters, while the other three are less obvious. To simplify, we just use the first one, hopefully to train a model to generate front facing characters by de-masking pixels step by step. Training training The goal to learn a process to convert fully noised/masked image to a front-facing character pixel art. If we know the ending pixel art, we can design a process like this: at each time step, a masked pixel can decide either stay masked or jump to the ending color, an un-masked pixel always stays in the ending color. The cumulative probability of jump from 0 to t is \ \alpha t\ , and it satisfies \ \alpha 0 = 0\ and \ \alpha 1 = 1\ . This \ \alpha t\ is called denoise schedule. When the ending pixel art is not given, the denoise schedule is still \ \alpha t\ , but the ending color to jump to is unknown. The best guess we make here is the most likely ending state given the current partially noised image and time t. We build a multi-head transformer to predict the ending state, the model can be further simplified to skip t as the input as the time information is already embedded in the number of masked pixels. img size = 16 p mask = 256 n vocab = 257 n seq = img size img size class Net nn.Module : def init self : super Net, self . init self.emb = nn.Embedding n vocab, dim self.pos emb = nn.Embedding n seq, dim self.blocks = nn.Sequential Block head size=dim // n head for in range n layer self.proj = nn.Linear dim, n vocab-1 def forward self, x : x, B:n seq x = self.emb x + self.pos emb torch.arange n seq, device=device B:n seq:dim x = self.blocks x return self.proj x B:n seq:n vocab-1 t = torch.rand batch size, device=device kappa = torch.bernoulli t.view batch size,1,1 .expand batch size, img size, img size .long .to device masked pos = kappa == 0 xt = p mask 1-kappa + images kappa z = model xt.view batch size,-1 loss = F.cross entropy z.view -1, n vocab-1 masked pos.view -1 , images.view -1 masked pos.view -1 Sampling sampling x = p mask torch.ones n sample, img size, img size , device=device .long alpha prev = 0 for step in range steps : alpha curr = step/steps alpha t = alpha curr - alpha prev / 1-alpha prev torch.ones n sample, device=device kappa = torch.bernoulli alpha t.view n sample,1,1 .expand n sample, img size, img size .long .to device kappa = kappa x == p mask .long logits = model x.view n sample, -1 sample = torch.multinomial F.softmax logits, dim=-1 .view n sample img size img size, -1 , num samples=1 .view n sample, img size, img size x = sample.view n sample, img size, img size kappa + x 1-kappa alpha prev = alpha curr