| 126 | |
| 127 | |
| 128 | class UNet(nn.Module): |
| 129 | def __init__(self, c_in=3, c_out=3, time_dim=256, device="cuda"): |
| 130 | super().__init__() |
| 131 | self.device = device |
| 132 | self.time_dim = time_dim |
| 133 | self.inc = DoubleConv(c_in, 64) |
| 134 | self.down1 = Down(64, 128) |
| 135 | self.sa1 = SelfAttention(128, 32) |
| 136 | self.down2 = Down(128, 256) |
| 137 | self.sa2 = SelfAttention(256, 16) |
| 138 | self.down3 = Down(256, 256) |
| 139 | self.sa3 = SelfAttention(256, 8) |
| 140 | |
| 141 | self.bot1 = DoubleConv(256, 512) |
| 142 | self.bot2 = DoubleConv(512, 512) |
| 143 | self.bot3 = DoubleConv(512, 256) |
| 144 | |
| 145 | self.up1 = Up(512, 128) |
| 146 | self.sa4 = SelfAttention(128, 16) |
| 147 | self.up2 = Up(256, 64) |
| 148 | self.sa5 = SelfAttention(64, 32) |
| 149 | self.up3 = Up(128, 64) |
| 150 | self.sa6 = SelfAttention(64, 64) |
| 151 | self.outc = nn.Conv2d(64, c_out, kernel_size=1) |
| 152 | |
| 153 | def pos_encoding(self, t, channels): |
| 154 | inv_freq = 1.0 / ( |
| 155 | 10000 |
| 156 | ** (torch.arange(0, channels, 2, device=self.device).float() / channels) |
| 157 | ) |
| 158 | pos_enc_a = torch.sin(t.repeat(1, channels // 2) * inv_freq) |
| 159 | pos_enc_b = torch.cos(t.repeat(1, channels // 2) * inv_freq) |
| 160 | pos_enc = torch.cat([pos_enc_a, pos_enc_b], dim=-1) |
| 161 | return pos_enc |
| 162 | |
| 163 | def forward(self, x, t): |
| 164 | t = t.unsqueeze(-1).type(torch.float) |
| 165 | t = self.pos_encoding(t, self.time_dim) |
| 166 | |
| 167 | x1 = self.inc(x) |
| 168 | x2 = self.down1(x1, t) |
| 169 | x2 = self.sa1(x2) |
| 170 | x3 = self.down2(x2, t) |
| 171 | x3 = self.sa2(x3) |
| 172 | x4 = self.down3(x3, t) |
| 173 | x4 = self.sa3(x4) |
| 174 | |
| 175 | x4 = self.bot1(x4) |
| 176 | x4 = self.bot2(x4) |
| 177 | x4 = self.bot3(x4) |
| 178 | |
| 179 | x = self.up1(x4, x3, t) |
| 180 | x = self.sa4(x) |
| 181 | x = self.up2(x, x2, t) |
| 182 | x = self.sa5(x) |
| 183 | x = self.up3(x, x1, t) |
| 184 | x = self.sa6(x) |
| 185 | output = self.outc(x) |