model card: link to the write-up

beb26e4 verified about 2 months ago

2.62 kB

license: cc-by-nc-4.0
pipeline_tag: text-to-image
library_name: pytorch
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

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.

Files

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

Usage

Clone the GitHub repo, fetch the frozen FLUX.2 VAE (gated, needs your own HF token; it is not redistributed here) and flan-t5-base, then:

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

License

Weights are released under CC BY-NC 4.0. They depend at inference on the FLUX.2 autoencoder (FLUX.2-dev non-commercial licence, gated) and were trained in part on images generated by FLUX.1-dev (FLUX-Reason-6M); use is non-commercial accordingly.