---
license: mit
pipeline_tag: text-to-image
tags:
- model_hub_mixin
- pytorch_model_hub_mixin
datasets:
- Artificio/WikiArt
---

This model has been pushed to the Hub using the [PytorchModelHubMixin](https://huggingface.co/docs/huggingface_hub/package_reference/mixins#huggingface_hub.PyTorchModelHubMixin) integration:
- Library: https://hf-mirror.com/kaupane/DiT-Wikiart-Small

### Model Description
This model is a DiT (diffusion transformer) trained on Wikiart dataset https://huggingface.co/datasets/Artificio/WikiArt from scratch. It is designed to generate art images given art genre and art style.

### Model Architecture
The model largely mirrors classic DiT architecture described in the paper *Scalable Diffusion Models with Transformers* with slight modifications:
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.

The model has three variants:
- 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
To use the model, install the "huggingface_hub" library and download modeling_dit_wikiart.py in "Files and versions" for model definition. After that you can use the model using the following code:
```python
from modeling_dit_wikiart import DiTWikiartModel

model = DiTWikiartModel.from_pretrained("kaupane/DiT-Wikiart-Small")
num_samples = 8
noisy_latents = torch.randn(num_samples,4,32,32)
predicted_noise = model(noisy_latents)
print(predicted_noise)
```
The model is paired with stabilityai/sd-vae-ft-ema.