cortex-image-cat

A small text-to-image diffusion model trained from scratch on cat photos at 64×64. No pretrained backbone, no CLIP, no VAE: the whole stack (BPE tokenizer, transformer text encoder, UNet, diffusion process, sampler) is implemented in the cortex-image project and fits in ~20.6 M parameters.

This release is the checkpoint of the 12,000-step Kaggle run (2×T4, DDP, fp16), selected as the first run that produces recognisable cats rather than colour blobs.

Generated through the project's own HTTP API with the default settings (40 DDIM steps, guidance 3.0): a photo of a cat, a photo of a black cat, a photo of a orange cat, a photo of a white cat.

What it can and cannot do

It works: the prompt reliably controls coat colour and tone — black, white, gray, orange, brown and cream come out in the right hue and the right order of brightness. Silhouette, ears, fur texture and eyes are recognisable on at least a few seeds per batch.

It does not: produce photo-realistic or reliably well-formed animals. At 64×64 with this amount of training you get soft, low-detail cats, and individual seeds still fail. There is no compositional or spatial control ("a cat on a sofa" is not understood), because the text conditioning is far too weak for that — see Training data below.

Sampler settings matter

guidance	steps	result
3.0	40	recommended: cleanest cats. This is the project default
1.0	40–100	dark, muddy, low contrast
6.0	40–100	oversaturated, breaks down into colour blobs

Top row: real training images. Then guidance 1.0 / 3.0 / 6.0 at 40 steps, then 1.0 / 3.0 / 6.0 at 100 steps, all with the prompt a photo of a cat and the same seed.

img2img (image + text → image) works and is safe at low strength: the sampler preserves the structure of the input image, so strength ≤ 0.3 is effectively a colour/texture restyle.

Training progress

Loss 0.4735 → 0.0548 over 12,000 steps (batch 128 effective, ~1.53 M images seen, ~34 passes over the dataset), at ~200 images/s on 2×T4. The curve flattens hard after step ~10 000: mean loss over the last 2 000 steps is 0.0586 against 0.1351 over the first 1 200, so most of the visible improvement happens early.

Training data

The Kaggle cats-vs-dogs archive (Microsoft Research), Cat class only: 12,498 images resized to 64×64.

Important caveat: that archive ships only a binary cat/dog label, which is far too little signal for text conditioning. Instead of a captioning model, this model was trained on pseudo-captions derived from image statistics — the median hue/saturation/value of the centre crop is bucketed into a coat word (black, white, gray, orange, brown, cream) and expanded into templates such as a photo of a {coat} cat. That is why colour is obeyed well and anything more complex is not: the model has never seen a real sentence describing an image.

Usage

The checkpoint carries its own config and tokenizer, but loading it needs the cortex-image package (custom UNet + diffusion code, not diffusers):

pip install torch pillow
git clone <this project>          # provides the cortex_image package
cd cortex-image

from cortex_image.pipeline import load_checkpoint, generate_batch
from cortex_image.viz import tensor_to_pil

loaded = load_checkpoint("model.pt", device="cpu", use_ema=True)
print(loaded.step, loaded.image_size, loaded.model.cfg.base_channels)

tensor, seed = generate_batch(
    loaded,
    ["a photo of a cat", "a photo of a black cat", "a photo of a orange cat"],
    n_per_prompt=2,
    steps=40,           # 40 is plenty; 60 is the UI default
    guidance=3.0,       # keep it at 3.0
    seed=1234,
)
for i in range(tensor.shape[0]):
    tensor_to_pil(tensor[i:i + 1]).save(f"out_{i:02d}.png")   # 64×64 PNG

Expect ~2 s per image for a batch of 4 at 40 steps on an Apple-silicon MPS device.

Or use the project's own web UI over this checkpoint:

python3 -m cortex_image.webui --ckpt model.pt --port 7860

Files

file	what it is
model.pt	full checkpoint: model weights, EMA weights, optimizer state, tokenizer and config, at step 12000
config.json	the training config used for this run (readable copy of what is embedded in model.pt)
samples/	generated samples and the real/generated comparison used above

model.pt is the complete resumable state, so it can also be used to continue training (--resume model.pt). It stores the EMA alongside the live weights; inference uses the EMA (use_ema=True), and this release was verified with the EMA weights.

Architecture

parameters	20.61 M
resolution	64×64, RGB
UNet	base_channels 64, mults [1,2,4], 2 res blocks per level, self-attention at 16×16, 4 heads
text encoder	4 transformer layers, 256 dim, 4 heads, 24 max tokens, BPE vocab 323
diffusion	cosine schedule, 1000 timesteps, v-prediction, min-SNR γ = 5.0, conditioning dropout 0.1 (enables CFG)
sampler	deterministic DDIM (40–100 steps), classifier-free guidance, optional img2img init
optimizer	AdamW, lr 2e-4 with 500-step warmup, grad clip 1.0, EMA decay 0.9995

Intended use and limitations

This is a research and learning artefact for a from-scratch diffusion implementation, at 64×64. It is not a product model and should not be used to generate depictions of real people, nor for any purpose requiring reliable or photorealistic output. No safety filtering, no watermarking, no provenance metadata.

Because the captions were derived from colour statistics rather than human-written descriptions, the model inherits any bias in that bucketing (a "black cat" bucket is defined by centre-crop brightness, so dark-furred cats in bright rooms are mislabelled). Cat photos dominate the Cat class of the source archive, so breed and pose coverage is whatever that archive contains.

Licensing

The model weights are released for research use. They were trained on the Microsoft Research cats-vs-dogs dataset, whose own terms continue to apply to anything derived from it; the image files themselves are not redistributed here. No licence is asserted beyond that.

Citation

@misc{cortex-image-cat,
  title  = {cortex-image-cat: a 20M-parameter from-scratch text-to-image diffusion model at 64x64},
  author = {Ilides},
  year   = {2026},
  note   = {Trained on the Kaggle cats-vs-dogs archive with pseudo-captions from coat-colour statistics}
}