Class-Conditioned UNet for Fashion MNIST

A diffusion model trained from scratch to generate class-conditioned images from the Fashion MNIST dataset. Built on top of diffusers' UNet2DModel, it accepts a class label (0–9) at inference time and generates the corresponding clothing category.

---

Model Architecture

The model wraps UNet2DModel with a learned class embedding that is concatenated to the noisy input as extra channels — a simple but effective conditioning strategy that avoids cross-attention overhead.

Class Labels

---

How to Use

1. Define ClassConditionedUnet

This wrapper must be defined before loading the model, as the pretrained weights cover only the inner UNet2DModel.

import torch
import torch.nn as nn
from diffusers import UNet2DModel

class ClassConditionedUnet(nn.Module):
    def __init__(self, num_classes=10, class_emb_size=4):
        super().__init__()
        self.class_emb = nn.Embedding(num_classes, class_emb_size)
        self.model = UNet2DModel(
            sample_size=28,
            in_channels=1 + class_emb_size,  # image + class embedding
            out_channels=1,
            layers_per_block=2,
            block_out_channels=(64, 128, 256),
            down_block_types=("DownBlock2D", "AttnDownBlock2D", "AttnDownBlock2D"),
            up_block_types=("AttnUpBlock2D", "AttnUpBlock2D", "UpBlock2D"),
        )

    def forward(self, x, t, class_labels):
        bs, ch, w, h = x.shape
        # Embed class labels and broadcast to spatial dimensions
        class_cond = self.class_emb(class_labels)                              # (bs, emb_size)
        class_cond = class_cond.view(bs, -1, 1, 1).expand(bs, -1, w, h)       # (bs, emb_size, w, h)
        net_input = torch.cat((x, class_cond), dim=1)                          # (bs, 1+emb_size, w, h)
        return self.model(net_input, t).sample

2. Load the Model

device = "cuda" if torch.cuda.is_available() else "cpu"

net = ClassConditionedUnet(num_classes=10, class_emb_size=4)
net.model = UNet2DModel.from_pretrained(
    "andreagemelli/UNet2DModel-fashion_mnist",
    use_safetensors=True,
)
net = net.to(device)
net.eval()

3. Generate Images

The snippet below generates 8 samples per class (80 images total) and displays them in a grid.

import torchvision
from diffusers import DDPMScheduler
from matplotlib import pyplot as plt
from tqdm.auto import tqdm

# 8 samples per class, all 10 classes → 80 images
n_per_class = 8
x = torch.randn(n_per_class * 10, 1, 28, 28).to(device)
y = torch.tensor([[i] * n_per_class for i in range(10)]).flatten().to(device)

# Scheduler
noise_scheduler = DDPMScheduler(num_train_timesteps=1000, beta_schedule="squaredcos_cap_v2")
noise_scheduler.set_timesteps(1000)

# Reverse diffusion loop
for t in tqdm(noise_scheduler.timesteps, desc="Sampling"):
    with torch.no_grad():
        residual = net(x, t, y)
    x = noise_scheduler.step(residual, t, x).prev_sample

# Visualize
grid = torchvision.utils.make_grid(x.detach().cpu().clip(-1, 1), nrow=n_per_class, normalize=True)
fig, ax = plt.subplots(figsize=(12, 6))
ax.imshow(grid.permute(1, 2, 0), cmap="Greys")
ax.set_title("Generated Fashion MNIST Images (rows = classes 0–9)")
ax.axis("off")
plt.tight_layout()
plt.show()

---

Training Details

---

Intended Use & Limitations

This model is intended for educational purposes — specifically to demonstrate class-conditioned image generation using the diffusers library.

Limitations:

• Output quality reflects the simplicity of Fashion MNIST (28×28 grayscale); this is not a high-fidelity generative model.
• Trained for only 5 epochs — longer training would improve sample quality and class coherence.
• The class embedding approach is intentionally minimal. More sophisticated conditioning (e.g. cross-attention or classifier-free guidance) would yield sharper class separation.

---

Credits

Training inspired by the Hugging Face Diffusion Models Course.