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
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
|
"""
UNet model for MNIST diffusion.
"""
import torch
import torch.nn as nn
import torch.nn.functional as F
import sys
sys.path.append("..")
from src.config import IMG_SIZE, NB_CHANNEL, NB_LABEL, DIFFU_STEPS
class SelfAttention(nn.Module):
"""Self-attention module for UNet."""
def __init__(self, nb_channels: int, nb_heads: int):
super().__init__()
self.attention = nn.MultiheadAttention(nb_channels, nb_heads)
def forward(self, x: torch.Tensor) -> torch.Tensor:
_, c, w, h = x.shape
x = x.view(-1, c, w * h)
x = x.permute(2, 0, 1)
x, _ = self.attention(x, x, x)
x = x.permute(1, 2, 0)
x = x.view(-1, c, w, h)
return x
class UNetMNIST(nn.Module):
"""
UNet architecture for MNIST diffusion model.
Inputs:
xt: image at step t (NB_CHANNEL x IMG_SIZE x IMG_SIZE)
t: step number (1)
vec: one-hot vector of the label (NB_LABEL)
Output:
Predicted noise or denoised image (NB_CHANNEL x IMG_SIZE x IMG_SIZE)
"""
def __init__(self, use_attention: bool = False):
super().__init__()
self.use_attention = use_attention
# Encoder for timestep t
self.encodet = nn.Linear(1, IMG_SIZE * IMG_SIZE)
# Encoder for label vector
self.encodevec = nn.Linear(NB_LABEL, IMG_SIZE * IMG_SIZE)
# UNet encoder (2 extra channels for t and vec)
self.conv1 = nn.Conv2d(NB_CHANNEL + 2, 64, 3, padding=1)
self.conv2 = nn.Conv2d(64, 64, 3, padding=1)
self.maxpool1 = nn.MaxPool2d(2, 2)
self.conv3 = nn.Conv2d(64, 128, 3, padding=1)
self.conv4 = nn.Conv2d(128, 128, 3, padding=1)
self.maxpool2 = nn.MaxPool2d(2, 2)
# Bottleneck
self.conv5 = nn.Conv2d(128, 256, 3, padding=1)
self.conv6 = nn.Conv2d(256, 256, 3, padding=1)
# Optional attention layers
if use_attention:
self.att1 = SelfAttention(128, 8)
self.att2 = SelfAttention(256, 8)
self.att3 = SelfAttention(256, 8)
# UNet decoder
self.upconv1 = nn.ConvTranspose2d(256, 128, 2, stride=2)
self.conv7 = nn.Conv2d(256, 128, 3, padding=1)
self.conv8 = nn.Conv2d(128, 128, 3, padding=1)
self.upconv2 = nn.ConvTranspose2d(128, 64, 2, stride=2)
self.conv9 = nn.Conv2d(128, 64, 3, padding=1)
self.conv10 = nn.Conv2d(64, 64, 3, padding=1)
# Output layer
self.conv11 = nn.Conv2d(64, NB_CHANNEL, 3, padding=1)
def forward(self, xt: torch.Tensor, t: torch.Tensor, vec: torch.Tensor) -> torch.Tensor:
"""
Forward pass of UNet.
Args:
xt: Noisy image at timestep t [B, C, H, W]
t: Timestep [B, 1]
vec: Label one-hot vector [B, NB_LABEL]
Returns:
Predicted noise or denoised image [B, C, H, W]
"""
# Encode timestep and label
t_enc = F.relu(self.encodet(t / DIFFU_STEPS))
t_enc = t_enc.view(-1, 1, IMG_SIZE, IMG_SIZE)
vec_enc = F.relu(self.encodevec(vec))
vec_enc = vec_enc.view(-1, 1, IMG_SIZE, IMG_SIZE)
# Concatenate input with embeddings
x = torch.cat((xt, t_enc, vec_enc), dim=1)
# Encoder path
x1 = F.relu(self.conv1(x))
x1 = F.relu(self.conv2(x1))
x2 = self.maxpool1(x1)
x2 = F.relu(self.conv3(x2))
x2 = F.relu(self.conv4(x2))
# Bottleneck
x3 = self.maxpool2(x2)
if self.use_attention:
x3 = self.att1(x3)
x3 = F.relu(self.conv5(x3))
if self.use_attention:
x3 = self.att2(x3)
x3 = F.relu(self.conv6(x3))
if self.use_attention:
x3 = self.att3(x3)
# Decoder path with skip connections
x4 = self.upconv1(x3)
x4 = torch.cat((x4, x2), dim=1)
x4 = F.relu(self.conv7(x4))
x4 = F.relu(self.conv8(x4))
x5 = self.upconv2(x4)
x5 = torch.cat((x5, x1), dim=1)
x5 = F.relu(self.conv9(x5))
x5 = F.relu(self.conv10(x5))
# Output
out = self.conv11(x5)
return out
|
