File size: 2,221 Bytes
0d83275 3f35e0d 0d83275 cdc1600 18d3148 0ad4f5f cdc1600 8769141 87316b7 1e1e5c0 283fc70 1e1e5c0 74511f3 1e1e5c0 87316b7 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 | ---
license: apache-2.0
datasets:
- mlfoundations/datacomp_1b
pipeline_tag: unconditional-image-generation
library_name: transformers
tags:
- imagegen
- unconditional-image
- custom_code
- tinyimagegen
---
# TinyImageGen
TinyImageGen is a small, fast unconditional image generation model. It features a total of 639k parameters and was trained on 50k images from DataComp for 15 epochs. While image models aren't our primary focus at FromZero, we decided to create our very first one.
## Architecture
TinyImageGen uses a custom architecture inspired by our text-to-text models, featuring mHC, Hadamard FFNs with SwiGLU intervals, 2D Axial RoPE, and a continuous diffusion objective.
- Hidden Size: `96`
- Hidden Layers: `6`
- Patch Size: `2×2`
- Attention Heads: `4`
- KV Heads: `2` (Grouped-Query Attention)
- Intermediate Size (for SwiGLU): `160`
- SwiGLU Interval: `3` (every 3rd layer)
- Number of Lanes: `4`
- RoPE Theta: `2500.0`
This architecture allows TinyImageGen to remain fast and parameter-efficient while still providing the effective depth and width of a much larger model.
## Training
As stated above, we trained TinyImageGen on 50k images from DataComp for 15 epochs.
- Final Loss: `0.2119`
## Hardware
- Ryzen 5 2600
## Generated Sample (x6 image grid)
<p align="left">
<img src="assets/sample_x6.png" width="15%" style="border-radius: 6px; margin-right: 10px;">
</p>
While the generated outputs are largely incoherent and unidentifiable, this is expected at such a small scale, and we make no claims otherwise.
## How to Use
```python
import torch
from transformers import AutoModel
from torchvision.utils import save_image
# Select device
device = "cuda" if torch.cuda.is_available() else "cpu"
# Load model directly from Hugging Face
model = AutoModel.from_pretrained(
"fromziro/TinyImageGen-0.6M",
trust_remote_code=True
).to(device)
# Generate 6 unconditional 32x32 images
with torch.no_grad():
samples = model.sample(num_samples=6, device=device, num_steps=50)
# Denormalize from [-1, 1] to [0, 1] and save grid
images = (samples * 0.5 + 0.5).clamp(0, 1)
save_image(images, "sample.png", nrow=3)
print("Saved samples to sample.png!")
```
## License
Apache 2.0. |