Add diffusers format: TinyDiTPipeline + TinyDiTTransformer2DModel (remote code), bundled T5 and FLUX.2 VAE
1c18331 verified | 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 | |
|  | |
|  | |
| ## 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. | |