---
library_name: cortex-image
pipeline_tag: text-to-image
license: other
tags:
  - text-to-image
  - diffusion
  - ddpm
  - ddim
  - from-scratch
  - cats
  - pytorch
---

# 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](https://github.com/) 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.

![prompts](samples/prompts.png)

*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 |

![guidance](samples/guidance.png)

*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

![progress](samples/progress.png)

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`):

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

```python
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:

```bash
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

```bibtex
@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}
}
```