(self, batch, **args)
| 546 | |
| 547 | @torch.no_grad() |
| 548 | def get_input_withmask(self, batch, **args): |
| 549 | out = super().get_input(batch, self.first_stage_key, **args) |
| 550 | mask = batch["mask"] |
| 551 | if len(mask.shape) == 3: |
| 552 | mask = mask[..., None] |
| 553 | mask = rearrange(mask, 'b h w c -> b c h w') |
| 554 | mask = mask.to(memory_format=torch.contiguous_format).float() |
| 555 | out += [mask] |
| 556 | |
| 557 | return out |
| 558 | |
| 559 | def training_step(self, batch, batch_idx): |
| 560 | if self.controller is not None: |