| 1318 | |
| 1319 | |
| 1320 | class DiffusionWrapper(pl.LightningModule): |
| 1321 | def __init__(self, diff_model_config, conditioning_key): |
| 1322 | super().__init__() |
| 1323 | self.sequential_cross_attn = diff_model_config.pop("sequential_crossattn", False) |
| 1324 | self.diffusion_model = instantiate_from_config(diff_model_config) |
| 1325 | self.conditioning_key = conditioning_key |
| 1326 | assert self.conditioning_key in [None, 'concat', 'crossattn', 'hybrid', 'adm', 'hybrid-adm', 'crossattn-adm'] |
| 1327 | |
| 1328 | def forward(self, x, t, c_concat: list = None, c_crossattn: list = None, c_adm=None): |
| 1329 | if self.conditioning_key is None: |
| 1330 | out = self.diffusion_model(x, t) |
| 1331 | elif self.conditioning_key == 'concat': |
| 1332 | xc = torch.cat([x] + c_concat, dim=1) |
| 1333 | out = self.diffusion_model(xc, t) |
| 1334 | elif self.conditioning_key == 'crossattn': |
| 1335 | if not self.sequential_cross_attn: |
| 1336 | cc = torch.cat(c_crossattn, 1) |
| 1337 | else: |
| 1338 | cc = c_crossattn |
| 1339 | out = self.diffusion_model(x, t, context=cc) |
| 1340 | elif self.conditioning_key == 'hybrid': |
| 1341 | xc = torch.cat([x] + c_concat, dim=1) |
| 1342 | cc = torch.cat(c_crossattn, 1) |
| 1343 | out = self.diffusion_model(xc, t, context=cc) |
| 1344 | elif self.conditioning_key == 'hybrid-adm': |
| 1345 | assert c_adm is not None |
| 1346 | xc = torch.cat([x] + c_concat, dim=1) |
| 1347 | cc = torch.cat(c_crossattn, 1) |
| 1348 | out = self.diffusion_model(xc, t, context=cc, y=c_adm) |
| 1349 | elif self.conditioning_key == 'crossattn-adm': |
| 1350 | assert c_adm is not None |
| 1351 | cc = torch.cat(c_crossattn, 1) |
| 1352 | out = self.diffusion_model(x, t, context=cc, y=c_adm) |
| 1353 | elif self.conditioning_key == 'adm': |
| 1354 | cc = c_crossattn[0] |
| 1355 | out = self.diffusion_model(x, t, y=cc) |
| 1356 | else: |
| 1357 | raise NotImplementedError() |
| 1358 | |
| 1359 | return out |
| 1360 | |
| 1361 | |
| 1362 | class LatentUpscaleDiffusion(LatentDiffusion): |