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

Usage

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

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

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.