---
license: cc-by-nc-4.0
pipeline_tag: text-to-image
library_name: diffusers
tags:
- diffusion-transformer
- rectified-flow
- text-to-image
- from-scratch
base_model:
- black-forest-labs/FLUX.2-dev
- google/flan-t5-base
datasets:
- LucasFang/FLUX-Reason-6M
- bghira/photo-concept-bucket
- detection-datasets/coco
---

# tinydit-256

A 210M-parameter text-to-image diffusion transformer trained from scratch on one GPU (3.5 days, 400k steps,
4.2M images at 256² area with aspect-ratio buckets). Frozen FLUX.2 autoencoder (32-channel latent) and frozen
flan-t5-base text encoder; the transformer, recipe, data pipeline and evaluation are original.

- 📝 Write-up: https://huggingface.co/blog/ivanmikhnenkov/tinydit-text-to-image-from-scratch-one-gpu
- 🎨 Try it in the browser: https://huggingface.co/spaces/ivanmikhnenkov/tinydit
- 💻 Code, figures and every design decision with sources: https://github.com/ivanmikhnenkov/tinydit (tag `v1-pretrain`)
- 🌐 Author: https://ivanmikhnenkov.com

![samples](grid_objects.jpg)

![same prompt over training](progress.jpg)

## Usage

With 🧨 diffusers (0.36 or newer); the transformer and pipeline are custom code in this repo, hence
`trust_remote_code`:

```python
import torch
from diffusers import DiffusionPipeline

pipe = DiffusionPipeline.from_pretrained("ivanmikhnenkov/tinydit-256", trust_remote_code=True,
                                         dtype=torch.bfloat16).to("cuda")
image = pipe("a lighthouse on a rocky coast at sunset", width=320, height=208,
             generator=torch.Generator("cuda").manual_seed(1)).images[0]
image.save("lighthouse.png")
```

- **Sizes**: trained on 256×256, 288×224, 224×288, 320×208 and 208×320 (W×H). Any multiple of 16 runs
  (2D RoPE); quality drops away from that area.
- **Defaults** are the training-time sampler: 20 Euler steps, CFG 4 against the empty prompt, timestep shift 2.8
  (in the scheduler config). 50 steps is slightly better; the shift is worth more than doubling the steps.
- Runs on CPU in fp32 too: about 40 s per 256² image at 20 steps on 4 ARM cores.
- The pipeline output matches the original sampler (`tinydit.sample`) to within 1e-5 in latents and 1/255 in
  pixels; `hub/parity.py` in the GitHub repo checks it.
- For fine-tuning: the transformer takes diffusers' flow-matching convention (`timestep` = σ·1000, predicts
  noise − data, as SD3 does), so the standard scheduler and training loops apply; gradient checkpointing works.

## Results (held-out prompts, 20 Euler steps, CFG 4, shift 2.8)

| FID | FD-DINOv2 | object accuracy (COCO classes, Faster R-CNN) | CLIP | PickScore | HPSv2.1 |
|---|---|---|---|---|---|
| 27.0 | 218 | 90% | 0.322 | 20.9 | 0.254 |

Strengths: single objects, animals, scenes, colours and simple relations, five aspect ratios. Known limits: readable
text, close faces, crowds, counts above three.

Architecture: adaLN-single DiT, dim 896 × 16 blocks, 2D RoPE, QK-norm, SwiGLU, 16 register tokens and 2 learned
null key/value slots per cross-attention; rectified flow with a 2.8 timestep shift, cosine and dispersive
auxiliary losses.

## Files

diffusers layout:

- `transformer/`: `TinyDiTTransformer2DModel` (`transformer_tinydit.py`), EMA weights (bf16), step 400k; its config
  holds the per-channel latent mean/std the latent space is whitened with
- `text_encoder/`, `tokenizer/`: flan-t5-base encoder
- `vae/`: the FLUX.2 autoencoder, as published in `black-forest-labs/FLUX.2-dev`
- `scheduler/`: `FlowMatchEulerDiscreteScheduler`, shift 2.8
- `pipeline_tinydit.py`: `TinyDiTPipeline`

Original layout, for the code on GitHub (same weights):

- `model.safetensors`: EMA weights (bf16), step 400k
- `config.json`: architecture and sampler defaults
- `stats.json`: per-channel latent mean/std; latents are whitened with these after encoding and un-whitened before decoding

```bash
python -m tinydit.sample --ckpt model.safetensors --config run1 --sets NOVEL=prompts/novel.txt --out out.png --shape 320x208
```

or run the local playground (`python -m tinydit.playground`) for interactive prompts, sampling trajectories and
attention maps.

## License

Weights are released under CC BY-NC 4.0. The bundled FLUX.2 autoencoder is released by Black Forest Labs under
Apache 2.0, as is flan-t5-base. The transformer was trained in part on images generated by FLUX.1-dev
(FLUX-Reason-6M); use is non-commercial accordingly.