README.md

---
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.