DiT_VAE/vae/utils/callbacks.py

import logging
import time, os
import numpy as np
from PIL import Image
from einops import rearrange
from omegaconf import OmegaConf

import torch
import torchvision
import pytorch_lightning as pl
from pytorch_lightning.callbacks import Callback
from pytorch_lightning.utilities.distributed import rank_zero_only
from pytorch_lightning.callbacks import Callback
from pytorch_lightning.utilities import rank_zero_info
from pytorch_lightning.callbacks import ModelCheckpoint, Callback, LearningRateMonitor
# ---------------------------------------------------------------------------------
class ImageLogger(Callback):
    def __init__(self, batch_frequency, max_images, clamp=True, increase_log_steps=True,
                 rescale=True, disabled=False, log_on_batch_idx=False, log_first_step=False,
                 fps=10, log_to_tblogger=True,
                 ):
        super().__init__()
        self.rescale = rescale
        self.batch_freq = batch_frequency
        self.max_images = max_images
        if int((pl.__version__).split('.')[1])>=7:
            self.logger_log_images = {
                pl.loggers.CSVLogger: self._testtube,
            }
        else:
            self.logger_log_images = {
                pl.loggers.TestTubeLogger: self._testtube,
            }
        self.log_steps = [2 ** n for n in range(int(np.log2(self.batch_freq)) + 1)]
        if not increase_log_steps:
            self.log_steps = [self.batch_freq]
        self.clamp = clamp
        self.disabled = disabled
        self.log_on_batch_idx = log_on_batch_idx
        self.log_first_step = log_first_step
        self.log_to_tblogger = log_to_tblogger
        self.save_fps = fps


    @rank_zero_only
    def _testtube(self, pl_module, images, batch_idx, split):
        """ log images and videos to tensorboard """

        for k in images:
            tag = f"{split}/{k}"
            if images[k].dim() == 5:
                video = images[k]
                n = video.shape[0]
                video = video.permute(2, 0, 1, 3, 4)
                frame_grids = [torchvision.utils.make_grid(framesheet, nrow=int(np.sqrt(n))) for framesheet in video]
                grid = torch.stack(frame_grids, dim=0)
                grid = (grid + 1.0) / 2.0
                grid = grid.unsqueeze(dim=0)
                pl_module.logger.experiment.add_video(
                    tag, grid,
                    global_step=pl_module.global_step)
            else:
                grid = torchvision.utils.make_grid(images[k])
                grid = (grid + 1.0) / 2.0
                pl_module.logger.experiment.add_image(
                    tag, grid,
                    global_step=pl_module.global_step)

    @rank_zero_only
    def log_local(self, save_dir, split, images,
                  global_step, current_epoch, batch_idx, rank_idx):
        """ save images and videos from images dict """

        root = os.path.join(save_dir, "images", split)
        os.makedirs(root, exist_ok=True)

        def save_img_grid(grid, path, rescale):
            if rescale:
                    grid = (grid + 1.0) / 2.0
            grid = grid.transpose(0, 1).transpose(1, 2).squeeze(-1)
            grid = grid.numpy()
            grid = (grid * 255).astype(np.uint8)
            os.makedirs(os.path.split(path)[0], exist_ok=True)
            Image.fromarray(grid).save(path)
        fps = images.pop('fps', None)
        fs = images.pop('frame_stride', None)

        for k in images:
            img = images[k]
            if isinstance(img, list) and isinstance(img[0], str):
                # a batch of captions
                filename = "string_{}_gs-{:06}_e-{:06}_b-{:06}_r-{:02}.txt".format(
                    k,
                    global_step,
                    current_epoch,
                    batch_idx,
                    rank_idx,
                    )
                path = os.path.join(root, filename)
                with open(path, 'w') as f:
                    for i, txt in enumerate(img):
                        f.write(f'idx={i}, txt={txt}\n')
                f.close()
            elif img.dim() == 5:
                # save video grids
                video = img # b,c,t,h,w
                n = video.shape[0]
                video = video.permute(2, 0, 1, 3, 4) # t,n,c,h,w
                frame_grids = [torchvision.utils.make_grid(framesheet, nrow=int(np.sqrt(n))) for framesheet in video] # [3, grid_h, grid_w]
                grid = torch.stack(frame_grids, dim=0) # stack in temporal dim [T, 3, grid_h, grid_w]
                if self.rescale:
                    grid = (grid + 1.0) / 2.0
                grid = (grid * 255).to(torch.uint8).permute(0, 2, 3, 1) # [T, 3, grid_h, grid_w] -> [T, grid_h, grid_w, 3]
                filename = "video_{}_gs-{:06}_e-{:06}_b-{:06}_r-{:02}.mp4".format(
                    k,
                    global_step,
                    current_epoch,
                    batch_idx,
                    rank_idx,
                    )
                filename = filename.split('.mp4')[0] + f'_fps{fps}.mp4' if fps is not None else filename
                filename = filename.split('.mp4')[0] + f'_fs{fs}.mp4' if fs is not None else filename
                path = os.path.join(root, filename)
                print('Save video ...')
                torchvision.io.write_video(path, grid, fps=self.save_fps, video_codec='h264', options={'crf': '10'})
                print('Finish!')
                
                # save frame sheet
                video_frames = rearrange(img, 'b c t h w -> (b t) c h w')
                t = img.shape[2]
                grid = torchvision.utils.make_grid(video_frames, nrow=t)
                filename = "framesheet_{}_gs-{:06}_e-{:06}_b-{:06}_r-{:02}.jpg".format(
                    k,
                    global_step,
                    current_epoch,
                    batch_idx,
                    rank_idx,
                    )
                path = os.path.join(root, filename)
                print('Save framesheet ...')
                save_img_grid(grid, path, self.rescale)
                print('Finish!')
            else:
                grid = torchvision.utils.make_grid(img, nrow=4)
                filename = "{}_gs-{:06}_e-{:06}_b-{:06}_r-{:02}.jpg".format(
                    k,
                    global_step,
                    current_epoch,
                    batch_idx,
                    rank_idx,
                    )
                path = os.path.join(root, filename)
                print('Save image grid ...')
                save_img_grid(grid, path, self.rescale)
                print('Finish!')
                
    @rank_zero_only
    def log_img(self, pl_module, batch, batch_idx, split="train", rank=0):
        """ generate images, then save and log to tensorboard """
        
        check_idx = batch_idx if self.log_on_batch_idx else pl_module.global_step
        
        if (self.check_frequency(check_idx) and
                hasattr(pl_module, "log_images") and
                callable(pl_module.log_images) and
                self.max_images > 0):
            logger = type(pl_module.logger)

            is_train = pl_module.training
            if is_train:
                pl_module.eval()
                torch.cuda.empty_cache()

            with torch.no_grad():
                log_func = pl_module.log_videos if hasattr(pl_module, 'is_video') and pl_module.is_video else pl_module.log_images
                images = log_func(batch, split=split)
                torch.cuda.empty_cache()

            # process images
            for k in images:
                if hasattr(images[k], 'shape'):
                    N = min(images[k].shape[0], self.max_images)
                    images[k] = images[k][:N]
                if isinstance(images[k], torch.Tensor):
                    images[k] = images[k].detach().cpu()
                    if self.clamp:
                        images[k] = torch.clamp(images[k], -1., 1.)
            
            print("Log local ...")
            self.log_local(pl_module.logger.save_dir, split, images,
                           pl_module.global_step, pl_module.current_epoch, 
                           batch_idx, rank)
            if self.log_to_tblogger:
                print("Log images to logger ...")
                logger_log_images = self.logger_log_images.get(logger, lambda *args, **kwargs: None)
                logger_log_images(pl_module, images, pl_module.global_step, split)
                print('Finish!')

            if is_train:
                pl_module.train()

    def check_frequency(self, check_idx):
        if ((check_idx % self.batch_freq) == 0 or (check_idx in self.log_steps)) and (
                check_idx > 0 or self.log_first_step):
            try:
                self.log_steps.pop(0)
            except IndexError as e:
                print(e)
                pass
            return True
        return False

    def on_train_batch_end(self, trainer, pl_module, outputs, batch, batch_idx, dataloader_idx=None):
        if not self.disabled and (pl_module.global_step > 0 or self.log_first_step):
            self.log_img(pl_module, batch, batch_idx, split="train", rank=trainer.global_rank)

    def on_validation_batch_end(self, trainer, pl_module, outputs, batch, batch_idx, dataloader_idx=None):
        if not self.disabled and pl_module.global_step > 0:
            self.log_img(pl_module, batch, batch_idx, split="val", rank=trainer.global_rank)
        if hasattr(pl_module, 'calibrate_grad_norm'):
            if (pl_module.calibrate_grad_norm and batch_idx % 25 == 0) and batch_idx > 0:
                self.log_gradients(trainer, pl_module, batch_idx=batch_idx)

# ---------------------------------------------------------------------------------
class CUDACallback(Callback):
    def on_train_epoch_start(self, trainer, pl_module):
        # Reset the memory use counter
        # lightning update
        if int((pl.__version__).split('.')[1])>=7:
            gpu_index = trainer.strategy.root_device.index
        else:
            gpu_index = trainer.root_gpu
        torch.cuda.reset_peak_memory_stats(gpu_index)
        torch.cuda.synchronize(gpu_index)
        self.start_time = time.time()

    def on_train_epoch_end(self, trainer, pl_module):
        if int((pl.__version__).split('.')[1])>=7:
            gpu_index = trainer.strategy.root_device.index
        else:
            gpu_index = trainer.root_gpu
        torch.cuda.synchronize(gpu_index)
        max_memory = torch.cuda.max_memory_allocated(gpu_index) / 2 ** 20
        epoch_time = time.time() - self.start_time

        try:
            max_memory = trainer.training_type_plugin.reduce(max_memory)
            epoch_time = trainer.training_type_plugin.reduce(epoch_time)
            rank_zero_info(f"Average Epoch time: {epoch_time:.2f} seconds")
            rank_zero_info(f"Average Peak memory {max_memory:.2f}MiB")
        except AttributeError:
            pass


# ---------------------------------------------------------------------------------
# for lower lighting
class SetupCallback_low(Callback):
    def __init__(self, resume, now, logdir, ckptdir, cfgdir, config, lightning_config, auto_resume, save_ptl_log=False):
        super().__init__()
        self.resume = resume
        self.now = now
        self.logdir = logdir
        self.ckptdir = ckptdir
        self.cfgdir = cfgdir
        self.config = config
        self.lightning_config = lightning_config
        self.auto_resume = auto_resume
        self.save_ptl_log = save_ptl_log

    def on_keyboard_interrupt(self, trainer, pl_module):
        if trainer.global_rank == 0:
            print("Summoning checkpoint.")
            ckpt_path = os.path.join(self.ckptdir, "last_summoning.ckpt")
            trainer.save_checkpoint(ckpt_path)

        # for old version lightning
    def on_pretrain_routine_start(self, trainer, pl_module):
        if trainer.global_rank == 0:
            # Create logdirs and save configs
            os.makedirs(self.logdir, exist_ok=True)
            os.makedirs(self.ckptdir, exist_ok=True)
            os.makedirs(self.cfgdir, exist_ok=True)
            
            if self.save_ptl_log:
                logger = logging.getLogger("pytorch_lightning")
                logger.addHandler(logging.FileHandler(os.path.join(self.logdir, "ptl_log.log")))

            if "callbacks" in self.lightning_config:
                if 'metrics_over_trainsteps_checkpoint' in self.lightning_config['callbacks']:
                    os.makedirs(os.path.join(self.ckptdir, 'trainstep_checkpoints'), exist_ok=True)
            print("Project config")
            print(OmegaConf.to_yaml(self.config))
            OmegaConf.save(self.config,
                           os.path.join(self.cfgdir, "{}-project.yaml".format(self.now)))

            print("Lightning config")
            print(OmegaConf.to_yaml(self.lightning_config))
            OmegaConf.save(OmegaConf.create({"lightning": self.lightning_config}),
                           os.path.join(self.cfgdir, "{}-lightning.yaml".format(self.now)))

        else:
            pass

# for higher lighting
class SetupCallback_high(Callback):
    def __init__(self, resume, now, logdir, ckptdir, cfgdir, config, lightning_config, auto_resume, save_ptl_log=False):
        super().__init__()
        self.resume = resume
        self.now = now
        self.logdir = logdir
        self.ckptdir = ckptdir
        self.cfgdir = cfgdir
        self.config = config
        self.lightning_config = lightning_config
        self.auto_resume = auto_resume
        self.save_ptl_log = save_ptl_log

    def on_keyboard_interrupt(self, trainer, pl_module):
        if trainer.global_rank == 0:
            print("Summoning checkpoint.")
            ckpt_path = os.path.join(self.ckptdir, "last_summoning.ckpt")
            trainer.save_checkpoint(ckpt_path)

    # RuntimeError: The `Callback.on_pretrain_routine_start` hook was removed in v1.8. Please use `Callback.on_fit_start` instead.
    def on_fit_start(self, trainer, pl_module):
        if trainer.global_rank == 0:
            # Create logdirs and save configs
            os.makedirs(self.logdir, exist_ok=True)
            os.makedirs(self.ckptdir, exist_ok=True)
            os.makedirs(self.cfgdir, exist_ok=True)
            
            if self.save_ptl_log:
                logger = logging.getLogger("pytorch_lightning")
                logger.addHandler(logging.FileHandler(os.path.join(self.logdir, "ptl_log.log")))

            if "callbacks" in self.lightning_config:
                if 'metrics_over_trainsteps_checkpoint' in self.lightning_config['callbacks']:
                    os.makedirs(os.path.join(self.ckptdir, 'trainstep_checkpoints'), exist_ok=True)
            print("Project config")
            print(OmegaConf.to_yaml(self.config))
            OmegaConf.save(self.config,
                           os.path.join(self.cfgdir, "{}-project.yaml".format(self.now)))

            print("Lightning config")
            print(OmegaConf.to_yaml(self.lightning_config))
            OmegaConf.save(OmegaConf.create({"lightning": self.lightning_config}),
                           os.path.join(self.cfgdir, "{}-lightning.yaml".format(self.now)))