| 325 | |
| 326 | @torch.no_grad() |
| 327 | def get_input(self, batch, k, bs=None, *args, **kwargs): |
| 328 | x, c = super().get_input(batch, self.first_stage_key, *args, **kwargs) |
| 329 | control_mask = batch[self.control_key] |
| 330 | if bs is not None: |
| 331 | control_mask = control_mask[:bs] |
| 332 | control_mask = control_mask.to(self.device) |
| 333 | control_mask = einops.rearrange(control_mask, 'b h w c -> b c h w') |
| 334 | control_mask = control_mask.to(memory_format=torch.contiguous_format).float() |
| 335 | |
| 336 | control_image = (batch["jpg"] + 1.0) / 2.0 |
| 337 | if bs is not None: |
| 338 | control_image = control_image[:bs] |
| 339 | control_image = control_image.to(self.device) |
| 340 | control_image = einops.rearrange(control_image, 'b h w c -> b c h w') |
| 341 | control_image = control_image.to(memory_format=torch.contiguous_format).float() |
| 342 | |
| 343 | return x, dict(c_crossattn=[c], c_concat_mask=[control_mask], c_concat_image=[control_image]) |
| 344 | |
| 345 | def apply_model(self, x_noisy, t, cond, *args, **kwargs): |
| 346 | assert isinstance(cond, dict) |