T.I.T.S.2 β€” Terrible Imagegen TranSformer 2

A 93M-parameter text-to-image diffusion transformer trained from scratch on a single RTX 4060 Ti in about 32 hours. 256x256, latent space, rectified flow.

It is the sequel to T.I.T.S., a 16M-parameter model that could draw a bus and nothing else. This one draws buses, motorcycles, dogs, pizza, snowy forests and castles. Unlike its predecessor it is an actual transformer, which means the name is now only 25% false advertising.

prompts

Prompts, left to right, top then bottom: "a castle on a cliff at sunset", "a cat wearing sunglasses", "a steaming bowl of ramen with an egg", "a red sports car on a mountain road", "a snowy forest path in winter", "a dragon flying over mountains", "a cup of coffee on a wooden table", "an astronaut riding a horse". Guidance 4.5, 30 steps.

The snowy path would pass as stock photography. The cat's sunglasses have become eyes. The dragon is a floating hillside. The astronaut is a war crime.

What it does well, and what it does not

Prompt Result
Landscapes, forests, roads, skies Frequently convincing
Vehicles β€” buses, motorcycles, cars Recognizable, correct number of wheels most of the time
Food on plates Reliably appetizing
Dogs and cats Real faces now. Bodies negotiate
Humans Do not
Text in images Scribble, as is tradition
Dragons, castles from fiction, anime CC12M has photos, not fantasy. It guesses politely

Specs

Parameters 93.19M (DiT) + frozen CLIP ViT-L/14 text encoder + frozen SD VAE
Architecture 12 blocks, dim 576, 9 heads. Self-attention + cross-attention to text + MLP, adaLN-Zero conditioning
Latent space SD-VAE-ft-mse, 4x32x32 -> 256x256 images, 2x2 patches = 256 tokens
Objective Rectified flow matching, logit-normal timestep sampling (SD3-style)
Sampling Euler, 30 steps, classifier-free guidance (10% caption dropout during training)
Training data 490,784 images from CC12M, watermark-filtered, LLaVA-written captions
Training 24 epochs / 360,550 steps / batch 32 / AdamW 1e-4 / bf16 / EMA 0.9995 / grad checkpointing
Hardware 1x RTX 4060 Ti (8GB), ~32 hours, sharing the GPU with whatever else was running
Final loss 0.7372 train / 0.7343 validation
Inference Well under a second per image on the training card

Validation loss stayed at or below training loss for all 24 epochs, so it never memorized anything and was still improving when we stopped. It stopped because we got bored, not because it converged.

Usage

pip install -r requirements.txt
python sample2.py --checkpoint tits2_ema_fp16.safetensors \
    --prompt "a snowy forest path in winter" --num_images 4 --guidance 4.5 --out out.png

Guidance 4-5 is the sweet spot. Below 3 it wanders, above 7 it fries.

Training your own

python prepare_data2.py --out_dir data2 --num 600000   # streams CC12M, caches VAE latents
./run_training2.sh                                     # 24 epochs, ~32h on a 4060 Ti

prepare_data2.py never writes a full-size image to disk: it streams pixparse/cc12m-wds for image bytes, joins it against opendiffusionai/cc12m-cleaned for watermark-filtered rows with LLaVA captions, and stores 8KB latents. 490k images = 3.9GB. Resumes at tar-file granularity, survives truncated tars, and self-heals if the caption window drifts out of sync with the image stream. All three of those behaviours exist because all three failures happened.

The 1-bit experiment

Because someone asked for it:

1bit

That is the same model with every weight matrix crushed to sign(w) * per_row_absmean. 101 matrices, 136.8 MB of weights down to 8.7 MB, whole file down to 58 MB. The output is two rectangles of premium pastel fog.

This is the expected result and not a bug. BitNet-style 1-bit models are trained with the constraint so the network learns to live within it; quantizing a finished model to 1 bit afterwards throws away everything except which side of zero each weight was on. quantize.py and sample_quant.py are included if you want to try 2 or 4 bits, which would probably survive.

Training progression

2k Step 2,000 β€” it knows "red" means red
10k Step 10,000 β€” palettes and horizons
50k Step 50,000 β€” buses become buses
final Step 360,000 β€” dogs have faces

Rows in the previews: motorcycle, golden retriever on a couch, pizza, red double decker bus.

Files

  • tits2_ema_fp16.safetensors β€” 186 MB, EMA weights, what you want
  • tits2_1bit.safetensors β€” 58 MB, the fog
  • Full training and data pipeline in the .py files

Limitations

Trained on web photos, so it inherits their biases and knows nothing that is not photographed a lot on the internet. No humans, no text, no fiction. Intended for entertainment and for finding out what 93M parameters and one gaming GPU can do.

Citation

Still no.

Downloads last month

-

Downloads are not tracked for this model. How to track
Inference Providers NEW
This model isn't deployed by any Inference Provider. πŸ™‹ Ask for provider support

Dataset used to train M1n1A1/MiniAI-TITS2