Modelpublic

M1n1A1/MiniAI-TITS2

sourceHugging Facemitupdated 4d agoView on Hugging Face

0likes

Model Card

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.

[image]

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

bash

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

bash

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:

[image]

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


[image]	Step 2,000 — it knows "red" means red
[image]	Step 10,000 — palettes and horizons
[image]	Step 50,000 — buses become buses
[image]	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.