kaupane/

DiT-Wikiart-Small

• Library: https://hf-mirror.com/kaupane/DiT-Wikiart-Small

Model Description

Model Architecture

1. Replaced ImageNet classes embeddings with Wikiart genres and styles embeddings;
2. Used post-norm instead of pre-norm;
3. Omitted final linear layer;
4. Replaced sin-cos-2d positional embedding with learned positional embedding;
5. Models only predict noise and don't learn sigma;
6. Setting patch_size=2 for all model variants;
7. Models have different size settings. Please check modeling_dit_wikiart.py in this repository for more details if you are interested.

• S: small, num_blocks=8, hidden_size=384, num_heads=6, total_params=20M;
• B: base, num_blocks=12, hidden_size=640, num_heads=10, total_params=90M;
• L: large, num_blocks=16, hidden_size=896, num_heads=14, total_params=234M.

Training Procedure

• dataset: all model variants were trained on 103K Wikiart dataset with data augmentation by horizontal flipping.
• optimizer: AdamW with default settings.
• learning rate: linear warmup for first 1% steps where learning rate reached a maximum of 3e-4, then cosine decay to zero in following steps.
• epochs and batch size:
  • S: 96 epochs with batch size of 176,
  • B: 120 epochs with batch size of 192,
  • L: 144 epochs with batch size of 192
• device:
  • S: single RTX 4060ti 16G for 24 hrs,
  • B: single RTX 4060ti 16G for 90 hrs,
  • L: single RTX 4090D 24G for 48 hrs, followed by single RTX 4060ti 16G for 100 hrs.
• loss curve: all variants witnessed a dramatic loss in the first epoch from above 1.0000 to around 0.2000, followed by a much slower decrease to finally reach loss=0.1600 at 20th epoch. DiT-S finally reached 0.1590; DiT-B finally reached 0.1525; DiT-L finally reached 0.1510.

Performance and Limitations

• The models demonstrates basic abilities to understand genres and styles and produce visually-appealing paintings (at first glance).
• Limitations include:
  • Failure to understand complex structures like human faces, buildings, etc.
  • Occassional modal collapse when asked to generate genres or styles rarely seen in the dataset. style like minimalism and genre like uroshi-e for example.
  • Resolution limited to 256x256
  • Trained on Wikiart dataset, therefore unable to generate out-of-scope images

How to use it