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