(self, batch)
| 101 | os.makedirs(os.path.join(self.logdir, 'images_val'), exist_ok=True) |
| 102 | |
| 103 | def prepare_batch_data(self, batch): |
| 104 | # prepare stable diffusion input |
| 105 | cond_imgs = batch['cond_imgs'] # (B, C, H, W) |
| 106 | cond_imgs = cond_imgs.to(self.device) |
| 107 | |
| 108 | # random resize the condition image |
| 109 | cond_size = np.random.randint(128, 513) |
| 110 | cond_imgs = v2.functional.resize(cond_imgs, cond_size, interpolation=3, antialias=True).clamp(0, 1) |
| 111 | |
| 112 | target_imgs = batch['target_imgs'] # (B, 6, C, H, W) |
| 113 | target_imgs = v2.functional.resize(target_imgs, 320, interpolation=3, antialias=True).clamp(0, 1) |
| 114 | target_imgs = rearrange(target_imgs, 'b (x y) c h w -> b c (x h) (y w)', x=3, y=2) # (B, C, 3H, 2W) |
| 115 | target_imgs = target_imgs.to(self.device) |
| 116 | |
| 117 | return cond_imgs, target_imgs |
| 118 | |
| 119 | @torch.no_grad() |
| 120 | def forward_vision_encoder(self, images): |
no outgoing calls
no test coverage detected