| 27 | |
| 28 | @rank_zero_only |
| 29 | def log_local(self, save_dir, split, images, global_step, current_epoch, batch_idx): |
| 30 | root = os.path.join(save_dir, "image_log", split) |
| 31 | # root = os.path.join(save_dir, "image_log", split+"_titan") |
| 32 | path_image = {} |
| 33 | for k in images: |
| 34 | grid = torchvision.utils.make_grid(images[k], nrow=4) |
| 35 | if self.rescale: |
| 36 | grid = (grid + 1.0) / 2.0 # -1,1 -> 0,1; c,h,w |
| 37 | grid = grid.transpose(0, 1).transpose(1, 2).squeeze(-1) |
| 38 | grid = grid.numpy() |
| 39 | grid = (grid * 255).astype(np.uint8) |
| 40 | filename = "{}_gs-{:06}_e-{:06}_b-{:06}.png".format(k, global_step, current_epoch, batch_idx) |
| 41 | path = os.path.join(root, filename) |
| 42 | os.makedirs(os.path.split(path)[0], exist_ok=True) |
| 43 | grid = Image.fromarray(grid) |
| 44 | grid.save(path) |
| 45 | path_image[k] = grid |
| 46 | grid_control = path_image["control_mask"] |
| 47 | grid_samples = path_image["samples_cfg_scale_9.00_mask"] |
| 48 | blended_filename = "{}_gs-{:06}_e-{:06}_b-{:06}.png".format("blended", global_step, current_epoch, batch_idx) |
| 49 | blended_path = os.path.join(root, blended_filename) |
| 50 | blended_image = Image.blend(grid_control, grid_samples, alpha=0.5) |
| 51 | blended_image.save(blended_path) |
| 52 | |
| 53 | def log_img(self, pl_module, batch, batch_idx, split="train"): |
| 54 | check_idx = batch_idx # if self.log_on_batch_idx else pl_module.global_step |