(self)
| 229 | |
| 230 | class SDUNet(torch.nn.Module): |
| 231 | def __init__(self): |
| 232 | super().__init__() |
| 233 | self.time_proj = Timesteps(320) |
| 234 | self.time_embedding = torch.nn.Sequential( |
| 235 | torch.nn.Linear(320, 1280), |
| 236 | torch.nn.SiLU(), |
| 237 | torch.nn.Linear(1280, 1280) |
| 238 | ) |
| 239 | self.conv_in = torch.nn.Conv2d(4, 320, kernel_size=3, padding=1) |
| 240 | |
| 241 | self.blocks = torch.nn.ModuleList([ |
| 242 | # CrossAttnDownBlock2D |
| 243 | ResnetBlock(320, 320, 1280), |
| 244 | AttentionBlock(8, 40, 320, 1, 768, eps=1e-6), |
| 245 | PushBlock(), |
| 246 | ResnetBlock(320, 320, 1280), |
| 247 | AttentionBlock(8, 40, 320, 1, 768, eps=1e-6), |
| 248 | PushBlock(), |
| 249 | DownSampler(320), |
| 250 | PushBlock(), |
| 251 | # CrossAttnDownBlock2D |
| 252 | ResnetBlock(320, 640, 1280), |
| 253 | AttentionBlock(8, 80, 640, 1, 768, eps=1e-6), |
| 254 | PushBlock(), |
| 255 | ResnetBlock(640, 640, 1280), |
| 256 | AttentionBlock(8, 80, 640, 1, 768, eps=1e-6), |
| 257 | PushBlock(), |
| 258 | DownSampler(640), |
| 259 | PushBlock(), |
| 260 | # CrossAttnDownBlock2D |
| 261 | ResnetBlock(640, 1280, 1280), |
| 262 | AttentionBlock(8, 160, 1280, 1, 768, eps=1e-6), |
| 263 | PushBlock(), |
| 264 | ResnetBlock(1280, 1280, 1280), |
| 265 | AttentionBlock(8, 160, 1280, 1, 768, eps=1e-6), |
| 266 | PushBlock(), |
| 267 | DownSampler(1280), |
| 268 | PushBlock(), |
| 269 | # DownBlock2D |
| 270 | ResnetBlock(1280, 1280, 1280), |
| 271 | PushBlock(), |
| 272 | ResnetBlock(1280, 1280, 1280), |
| 273 | PushBlock(), |
| 274 | # UNetMidBlock2DCrossAttn |
| 275 | ResnetBlock(1280, 1280, 1280), |
| 276 | AttentionBlock(8, 160, 1280, 1, 768, eps=1e-6), |
| 277 | ResnetBlock(1280, 1280, 1280), |
| 278 | # UpBlock2D |
| 279 | PopBlock(), |
| 280 | ResnetBlock(2560, 1280, 1280), |
| 281 | PopBlock(), |
| 282 | ResnetBlock(2560, 1280, 1280), |
| 283 | PopBlock(), |
| 284 | ResnetBlock(2560, 1280, 1280), |
| 285 | UpSampler(1280), |
| 286 | # CrossAttnUpBlock2D |
| 287 | PopBlock(), |
| 288 | ResnetBlock(2560, 1280, 1280), |
nothing calls this directly
no test coverage detected