Nano-Catana

TLDR

I have got my hands on a couple of GPUs ranging from NVIDIA L40S to H100s. I have scraped a dataset of 100k+ cat pictures, annotated them using LLMs and trained a couple of diffusion models from scratch. The models are trained on 512x512 images and can generate pretty convincing cat images. The models are trained using Lightning AI and PyTorch. While really silly and not the most useful project, it was an awesome way to learn about image diffusion, training models from scratch and managing training runs across multiple GPU clusters. You can even try it yourself at research.kocicividea.cz, where the model nano catana ttc (text-to-cat) generates cats from your prompts.

The why?

Everyone uses image models these days, but I had only some vague idea how they work under the hood. So I have decided to train my own image model, completely from scratch. No fine-tuning of Stable Diffusion, no pretrained weights for the image part. And because the internet is made of cats, the choice of the dataset was obvious.

Thanks to my university, I have access to MetaCentrum, the Czech national grid, which has a bunch of GPU nodes with L40S, A100 and H100 cards. That is a lot more compute than my laptop, so there was no excuse not to try.

The dataset

I have scraped over 100k cat pictures from the internet. A lot of them were not great. Memes, screenshots, cat food packaging, blurry photos of something orange. To filter them and to get captions for the text-to-image model, I have run every image through a vision LLM (Qwen) that described the cat, its pose, the background and flagged images that were not useful. After the cleanup, around 47k annotated images remained. The whole dataset is public on Hugging Face.

The models

I have trained a couple of different models, each one teaching me something new:

1. cat-diffusion, a classic unconditional DDPM. You can't tell it what to generate, it just makes a random cat from a seed. This was my first model and the "hello world" of image diffusion.
2. flowmatch-cats, the same idea but trained with flow matching instead of the classic noise prediction. Flow matching is what newer models like Stable Diffusion 3 and Flux use, and it trains noticeably faster.
3. flowmatch-cats-tti, the text-to-image (or rather text-to-cat) model. It is a pixel-space UNet conditioned on CLIP text embeddings, so you can write something like "a fluffy grey cat sleeping on a sofa" and get, well, a fluffy grey cat sleeping on a sofa. Most of the time.

# one flow matching training step, simplified
noise = torch.randn_like(images)
t = torch.rand(images.shape[0], device=images.device)
x_t = (1 - t) * images + t * noise          # blend the image with noise
target = noise - images                      # the velocity the model should predict
pred = unet(x_t, t, encoder_hidden_states=clip(prompts))
loss = F.mse_loss(pred, target)

Simplified, the real training loop has a bit more going on.

Try it!!

A model sitting on a storage server is no fun, so I have made a website where anyone can generate cats: research.kocicividea.cz. The frontend is a NextJS app running on Cloudflare Workers. All generations are public, so you can see what other people are making. You can also browse the training dataset and even donate a photo of your own cat. An LLM checks that it really is a cat and writes a caption for it, so it can be used in future training runs.