src/trainer/marigold_xl_trainer.py

# An official reimplemented version of Marigold training script.
# Last modified: 2024-04-29
#
# Copyright 2023 Bingxin Ke, ETH Zurich. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# --------------------------------------------------------------------------
# If you find this code useful, we kindly ask you to cite our paper in your work.
# Please find bibtex at: https://github.com/prs-eth/Marigold#-citation
# If you use or adapt this code, please attribute to https://github.com/prs-eth/marigold.
# More information about the method can be found at https://marigoldmonodepth.github.io
# --------------------------------------------------------------------------


import logging
import os
import pdb
import shutil
from datetime import datetime
from typing import List, Union
import safetensors
import numpy as np
import torch
from diffusers import DDPMScheduler
from omegaconf import OmegaConf
from torch.nn import Conv2d
from torch.nn.parameter import Parameter
from torch.optim import Adam
from torch.optim.lr_scheduler import LambdaLR
from torch.utils.data import DataLoader
from tqdm import tqdm
from PIL import Image
# import torch.optim.lr_scheduler

from marigold.marigold_pipeline import MarigoldPipeline, MarigoldDepthOutput
from src.util import metric
from src.util.data_loader import skip_first_batches
from src.util.logging_util import tb_logger, eval_dic_to_text
from src.util.loss import get_loss
from src.util.lr_scheduler import IterExponential
from src.util.metric import MetricTracker
from src.util.multi_res_noise import multi_res_noise_like
from src.util.alignment import align_depth_least_square
from src.util.seeding import generate_seed_sequence
from accelerate import Accelerator
import random

class MarigoldXLTrainer:
    def __init__(
        self,
        cfg: OmegaConf,
        model: MarigoldPipeline,
        train_dataloader: DataLoader,
        device,
        base_ckpt_dir,
        out_dir_ckpt,
        out_dir_eval,
        out_dir_vis,
        accumulation_steps: int,
        separate_list: List = None,
        val_dataloaders: List[DataLoader] = None,
        vis_dataloaders: List[DataLoader] = None,
        timestep_method: str = 'unidiffuser'
    ):
        self.cfg: OmegaConf = cfg
        self.model: MarigoldPipeline = model
        self.device = device
        self.seed: Union[int, None] = (
            self.cfg.trainer.init_seed
        )  # used to generate seed sequence, set to `None` to train w/o seeding
        self.out_dir_ckpt = out_dir_ckpt
        self.out_dir_eval = out_dir_eval
        self.out_dir_vis = out_dir_vis
        self.train_loader: DataLoader = train_dataloader
        self.val_loaders: List[DataLoader] = val_dataloaders
        self.vis_loaders: List[DataLoader] = vis_dataloaders
        self.accumulation_steps: int = accumulation_steps
        self.separate_list = separate_list
        self.timestep_method = timestep_method
        # Adapt input layers
        # if 8 != self.model.unet.config["in_channels"]:
        #     self._replace_unet_conv_in()
        # if 8 != self.model.unet.config["out_channels"]:
        #     self._replace_unet_conv_out()

        self.prompt = ['a view of a city skyline from a bridge',
                       'a man and a woman sitting on a couch',
                       'a black car parked in a parking lot next to the water',
                       'Enchanted forest with glowing plants, fairies, and ancient castle.',
                       'Futuristic city with skyscrapers, neon lights, and hovering vehicles.',
                       'Fantasy mountain landscape with waterfalls, dragons, and mythical creatures.']
        # self.generator = torch.Generator('cuda:0').manual_seed(1024)

        # Encode empty text prompt
        # self.model.encode_empty_text()
        # self.empty_text_embed = self.model.empty_text_embed.detach().clone().to(device)

        self.model.unet.enable_xformers_memory_efficient_attention()

        # Trainability
        self.model.vae.requires_grad_(False)
        self.model.text_encoder.requires_grad_(False)
        # self.model.unet.requires_grad_(True)

        grad_part = filter(lambda p: p.requires_grad, self.model.unet.parameters())

        # Optimizer !should be defined after input layer is adapted
        lr = self.cfg.lr
        self.optimizer = Adam(grad_part, lr=lr)

        total_params = sum(p.numel() for p in self.model.unet.parameters())
        total_params_m = total_params / 1_000_000
        print(f"Total parameters: {total_params_m:.2f}M")
        trainable_params = sum(p.numel() for p in self.model.unet.parameters() if p.requires_grad)
        trainable_params_m = trainable_params / 1_000_000
        print(f"Trainable parameters: {trainable_params_m:.2f}M")

        # LR scheduler
        lr_func = IterExponential(
            total_iter_length=self.cfg.lr_scheduler.kwargs.total_iter,
            final_ratio=self.cfg.lr_scheduler.kwargs.final_ratio,
            warmup_steps=self.cfg.lr_scheduler.kwargs.warmup_steps,
        )
        self.lr_scheduler = LambdaLR(optimizer=self.optimizer, lr_lambda=lr_func)

        # Loss
        self.loss = get_loss(loss_name=self.cfg.loss.name, **self.cfg.loss.kwargs)

        # Training noise scheduler
        self.training_noise_scheduler: DDPMScheduler = DDPMScheduler.from_pretrained(
            os.path.join(
                cfg.trainer.training_noise_scheduler.pretrained_path,
                "scheduler",
            )
        )
        self.prediction_type = self.training_noise_scheduler.config.prediction_type
        assert (
            self.prediction_type == self.model.scheduler.config.prediction_type
        ), "Different prediction types"
        self.scheduler_timesteps = (
            self.training_noise_scheduler.config.num_train_timesteps
        )

        # Eval metrics
        self.metric_funcs = [getattr(metric, _met) for _met in cfg.eval.eval_metrics]
        self.train_metrics = MetricTracker(*["loss", 'rgb_loss', 'depth_loss'])
        self.val_metrics = MetricTracker(*[m.__name__ for m in self.metric_funcs])
        # main metric for best checkpoint saving
        self.main_val_metric = cfg.validation.main_val_metric
        self.main_val_metric_goal = cfg.validation.main_val_metric_goal
        assert (
            self.main_val_metric in cfg.eval.eval_metrics
        ), f"Main eval metric `{self.main_val_metric}` not found in evaluation metrics."
        self.best_metric = 1e8 if "minimize" == self.main_val_metric_goal else -1e8

        # Settings
        self.max_epoch = self.cfg.max_epoch
        self.max_iter = self.cfg.max_iter
        self.gradient_accumulation_steps = accumulation_steps
        self.gt_depth_type = self.cfg.gt_depth_type
        self.gt_mask_type = self.cfg.gt_mask_type
        self.save_period = self.cfg.trainer.save_period
        self.backup_period = self.cfg.trainer.backup_period
        self.val_period = self.cfg.trainer.validation_period
        self.vis_period = self.cfg.trainer.visualization_period

        # Multi-resolution noise
        self.apply_multi_res_noise = self.cfg.multi_res_noise is not None
        if self.apply_multi_res_noise:
            self.mr_noise_strength = self.cfg.multi_res_noise.strength
            self.annealed_mr_noise = self.cfg.multi_res_noise.annealed
            self.mr_noise_downscale_strategy = (
                self.cfg.multi_res_noise.downscale_strategy
            )

        # Internal variables
        self.epoch = 0
        self.n_batch_in_epoch = 0  # batch index in the epoch, used when resume training
        self.effective_iter = 0  # how many times optimizer.step() is called
        self.in_evaluation = False
        self.global_seed_sequence: List = []  # consistent global seed sequence, used to seed random generator, to ensure consistency when resuming

    def _replace_unet_conv_in(self):
        # replace the first layer to accept 8 in_channels
        _weight = self.model.unet.conv_in.weight.clone()  # [320, 4, 3, 3]
        _bias = self.model.unet.conv_in.bias.clone()  # [320]
        zero_weight = torch.zeros(_weight.shape).to(_weight.device)
        _weight = torch.cat([_weight, zero_weight], dim=1)
        # _weight = _weight.repeat((1, 2, 1, 1))  # Keep selected channel(s)
        # half the activation magnitude
        # _weight *= 0.5
        # new conv_in channel
        _n_convin_out_channel = self.model.unet.conv_in.out_channels
        _new_conv_in = Conv2d(
            8, _n_convin_out_channel, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1)
        )
        _new_conv_in.weight = Parameter(_weight)
        _new_conv_in.bias = Parameter(_bias)
        self.model.unet.conv_in = _new_conv_in
        logging.info("Unet conv_in layer is replaced")
        # replace config
        self.model.unet.config["in_channels"] = 8
        logging.info("Unet config is updated")
        return

    def _replace_unet_conv_out(self):
        # replace the first layer to accept 8 in_channels
        _weight = self.model.unet.conv_out.weight.clone()  # [8, 320, 3, 3]
        _bias = self.model.unet.conv_out.bias.clone()  # [320]
        _weight = _weight.repeat((2, 1, 1, 1))  # Keep selected channel(s)
        _bias = _bias.repeat((2))
        # half the activation magnitude

        # new conv_in channel
        _n_convin_out_channel = self.model.unet.conv_out.out_channels
        _new_conv_out = Conv2d(
            _n_convin_out_channel, 8, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1)
        )
        _new_conv_out.weight = Parameter(_weight)
        _new_conv_out.bias = Parameter(_bias)
        self.model.unet.conv_out = _new_conv_out
        logging.info("Unet conv_out layer is replaced")
        # replace config
        self.model.unet.config["out_channels"] = 8
        logging.info("Unet config is updated")
        return

    def parallel_train(self, t_end=None, accelerator=None):
        logging.info("Start training")

        self.model, self.optimizer, self.train_loader, self.lr_scheduler = accelerator.prepare(
            self.model, self.optimizer, self.train_loader, self.lr_scheduler
        )
        self.accelerator = accelerator
        if self.val_loaders is not None:
            for idx, loader in enumerate(self.val_loaders):
                self.val_loaders[idx] = accelerator.prepare(loader)

        if os.path.exists(os.path.join(self.out_dir_ckpt, 'latest')):
            accelerator.load_state(os.path.join(self.out_dir_ckpt, 'latest'))
            self.load_miscs(os.path.join(self.out_dir_ckpt, 'latest'))

        self.train_metrics.reset()
        accumulated_step = 0
        for epoch in range(self.epoch, self.max_epoch + 1):
            self.epoch = epoch
            logging.debug(f"epoch: {self.epoch}")

            # Skip previous batches when resume
            for batch in skip_first_batches(self.train_loader, self.n_batch_in_epoch):
                self.model.unet.train()

                # globally consistent random generators
                if self.seed is not None:
                    local_seed = self._get_next_seed()
                    rand_num_generator = torch.Generator(device=self.model.device)
                    rand_num_generator.manual_seed(local_seed)
                else:
                    rand_num_generator = None

                # >>> With gradient accumulation >>>

                # Get data
                rgb = batch["rgb_norm"].to(self.model.device)
                depth_gt_for_latent = batch[self.gt_depth_type].to(self.model.device)
                batch_size = rgb.shape[0]

                if self.gt_mask_type is not None:
                    valid_mask_for_latent = batch[self.gt_mask_type].to(self.model.device)
                    invalid_mask = ~valid_mask_for_latent
                    valid_mask_down = ~torch.max_pool2d(
                        invalid_mask.float(), 8, 8
                    ).bool()
                    valid_mask_down = valid_mask_down.repeat((1, 4, 1, 1))

                with torch.no_grad():
                    # Encode image
                    rgb_latent = self.model.encode_rgb(rgb)  # [B, 4, h, w]
                    # Encode GT depth
                    gt_depth_latent = self.encode_depth(
                        depth_gt_for_latent
                    )  # [B, 4, h, w]

                # Sample a random timestep for each image
                if self.cfg.loss.depth_factor == 1:
                    rgb_timesteps = torch.zeros(
                        (batch_size),
                        device=self.model.device
                    ).long()  # [B]
                    depth_timesteps = torch.randint(
                        0,
                        self.scheduler_timesteps,
                        (batch_size,),
                        device=self.model.device,
                        generator=rand_num_generator,
                    ).long()  # [B]
                elif self.timestep_method == 'unidiffuser':
                    rgb_timesteps = torch.randint(
                        0,
                        self.scheduler_timesteps,
                        (batch_size,),
                        device=self.model.device,
                        generator=rand_num_generator,
                    ).long()  # [B]
                    depth_timesteps = torch.randint(
                        0,
                        self.scheduler_timesteps,
                        (batch_size,),
                        device=self.model.device,
                        generator=rand_num_generator,
                    ).long()  # [B]
                elif self.timestep_method == 'partition':
                    rand_num = random.random()
                    if rand_num < 0.3333:
                        # joint generation
                        rgb_timesteps = torch.randint(
                            0,
                            self.scheduler_timesteps,
                            (batch_size,),
                            device=self.model.device,
                            generator=rand_num_generator,
                        ).long()  # [B]
                        depth_timesteps = rgb_timesteps
                    elif rand_num < 0.6666:
                        # image2depth generation
                        rgb_timesteps = torch.zeros(
                            (batch_size),
                            device=self.model.device
                        ).long()  # [B]
                        depth_timesteps = torch.randint(
                            0,
                            self.scheduler_timesteps,
                            (batch_size,),
                            device=self.model.device,
                            generator=rand_num_generator,
                        ).long()  # [B]
                    else:
                        # depth2image generation
                        rgb_timesteps = torch.randint(
                            0,
                            self.scheduler_timesteps,
                            (batch_size,),
                            device=self.model.device,
                            generator=rand_num_generator,
                        ).long()  # [B]
                        depth_timesteps = torch.zeros(
                            (batch_size),
                            device=self.model.device
                        ).long()  # [B]

                # Sample noise
                if self.apply_multi_res_noise:
                    rgb_strength = self.mr_noise_strength
                    if self.annealed_mr_noise:
                        # calculate strength depending on t
                        rgb_strength = rgb_strength * (rgb_timesteps / self.scheduler_timesteps)
                    rgb_noise = multi_res_noise_like(
                        rgb_latent,
                        strength=rgb_strength,
                        downscale_strategy=self.mr_noise_downscale_strategy,
                        generator=rand_num_generator,
                        device=self.model.device,
                    )

                    depth_strength = self.mr_noise_strength
                    if self.annealed_mr_noise:
                        # calculate strength depending on t
                        depth_strength = depth_strength * (depth_timesteps / self.scheduler_timesteps)
                    depth_noise = multi_res_noise_like(
                        gt_depth_latent,
                        strength=depth_strength,
                        downscale_strategy=self.mr_noise_downscale_strategy,
                        generator=rand_num_generator,
                        device=self.model.device,
                    )
                else:
                    rgb_noise = torch.randn(
                        rgb_latent.shape,
                        device=self.model.device,
                        generator=rand_num_generator,
                    )  # [B, 8, h, w]

                    depth_noise = torch.randn(
                        gt_depth_latent.shape,
                        device=self.model.device,
                        generator=rand_num_generator,
                    )  # [B, 8, h, w]
                # Add noise to the latents (diffusion forward process)

                noisy_rgb_latents = self.training_noise_scheduler.add_noise(
                    rgb_latent, rgb_noise, rgb_timesteps
                )  # [B, 4, h, w]
                noisy_depth_latents = self.training_noise_scheduler.add_noise(
                    gt_depth_latent, depth_noise, depth_timesteps
                )  # [B, 4, h, w]

                noisy_latents = torch.cat(
                    [noisy_rgb_latents, noisy_depth_latents], dim=1
                ).float()  # [B, 8, h, w]

                # Text embedding
                batch_text_embed = []
                batch_pooled_text_embed = []
                for p in batch['text']:
                    prompt_embed, pooled_prompt_embed = self.model.encode_text(p)
                    batch_text_embed.append(prompt_embed)
                    batch_pooled_text_embed.append(pooled_prompt_embed)
                batch_text_embed = torch.cat(batch_text_embed, dim=0)
                batch_pooled_text_embed = torch.cat(batch_pooled_text_embed, dim=0)
                # input_ids = {k:v.squeeze().to(self.model.device) for k,v in batch['text'].items()}
                # prompt_embed, pooled_prompt_embed = self.model.encode_text(batch['text'])
                # text_embed = self.empty_text_embed.to(device).repeat(
                #     (batch_size, 1, 1)
                # )  # [B, 77, 1024]
                # Predict the noise residual
                add_time_ids = self.model._get_add_time_ids(
                    (batch['rgb_int'].shape[-2], batch['rgb_int'].shape[-1]), (0, 0), (batch['rgb_int'].shape[-2], batch['rgb_int'].shape[-1]), dtype=batch_text_embed.dtype
                )
                pdb.set_trace()
                dtype = self.model.unet.dtype
                added_cond_kwargs = {"text_embeds": batch_pooled_text_embed.to(self.model.device).to(dtype), "time_ids": add_time_ids.to(self.model.device).to(dtype)}
                model_pred = self.model.unet(
                    noisy_latents.to(self.model.unet.dtype), rgb_timesteps, depth_timesteps, encoder_hidden_states=batch_text_embed.to(dtype),
                    added_cond_kwargs=added_cond_kwargs, separate_list=self.separate_list
                ).sample  # [B, 4, h, w]
                if torch.isnan(model_pred).any():
                    logging.warning("model_pred contains NaN.")

                # Get the target for loss depending on the prediction type
                if "sample" == self.prediction_type:
                    rgb_target = rgb_latent
                    depth_target = gt_depth_latent
                elif "epsilon" == self.prediction_type:
                    rgb_target = rgb_latent
                    depth_target = gt_depth_latent
                elif "v_prediction" == self.prediction_type:
                    rgb_target = self.training_noise_scheduler.get_velocity(
                        rgb_latent, rgb_noise, rgb_timesteps
                    )  # [B, 4, h, w]
                    depth_target = self.training_noise_scheduler.get_velocity(
                        gt_depth_latent, depth_noise, depth_timesteps
                    )  # [B, 4, h, w]
                else:
                    raise ValueError(f"Unknown prediction type {self.prediction_type}")
                # Masked latent loss
                with accelerator.accumulate(self.model):
                    if self.gt_mask_type is not None:
                        depth_loss = self.loss(
                            model_pred[:, 4:, :, :][valid_mask_down].float(),
                            depth_target[valid_mask_down].float(),
                        )
                    else:
                        depth_loss = self.cfg.loss.depth_factor * self.loss(model_pred[:, 4:, :, :].float(),depth_target.float())

                    rgb_loss = (1 - self.cfg.loss.depth_factor) * self.loss(model_pred[:, 0:4, :, :].float(), rgb_target.float())
                    if self.cfg.loss.depth_factor == 1:
                        loss = depth_loss
                    else:
                        loss = rgb_loss + depth_loss

                    self.train_metrics.update("loss", loss.item())
                    self.train_metrics.update("rgb_loss", rgb_loss.item())
                    self.train_metrics.update("depth_loss", depth_loss.item())
                    # loss = loss / self.gradient_accumulation_steps
                    accelerator.backward(loss)
                    self.optimizer.step()
                    self.optimizer.zero_grad()
                    # loss.backward()
                    self.n_batch_in_epoch += 1
                    # print(accelerator.process_index, self.lr_scheduler.get_last_lr())
                    self.lr_scheduler.step(self.effective_iter)

                if accelerator.sync_gradients:
                    accumulated_step += 1

                if accumulated_step >= self.gradient_accumulation_steps:
                    accumulated_step = 0
                    self.effective_iter += 1

                    if accelerator.is_main_process:
                        # Log to tensorboard
                        if self.effective_iter == 1:
                            generator = torch.Generator(self.model.device).manual_seed(1024)
                            img = self.model.generate_rgbd(self.prompt, num_inference_steps=50, generator=generator,
                                                           show_pbar=True)
                            for idx in range(len(self.prompt)):
                                tb_logger.writer.add_image(f'image/{self.prompt[idx]}', img[idx], self.effective_iter)

                        accumulated_loss = self.train_metrics.result()["loss"]
                        rgb_loss = self.train_metrics.result()["rgb_loss"]
                        depth_loss = self.train_metrics.result()["depth_loss"]
                        tb_logger.log_dic(
                            {
                                f"train/{k}": v
                                for k, v in self.train_metrics.result().items()
                            },
                            global_step=self.effective_iter,
                        )
                        tb_logger.writer.add_scalar(
                            "lr",
                            self.lr_scheduler.get_last_lr()[0],
                            global_step=self.effective_iter,
                        )
                        tb_logger.writer.add_scalar(
                            "n_batch_in_epoch",
                            self.n_batch_in_epoch,
                            global_step=self.effective_iter,
                        )
                        logging.info(
                            f"iter {self.effective_iter:5d} (epoch {epoch:2d}): loss={accumulated_loss:.5f}, rgb_loss={rgb_loss:.5f}, depth_loss={depth_loss:.5f}"
                        )
                    accelerator.wait_for_everyone()

                    if self.save_period > 0 and 0 == self.effective_iter % self.save_period:
                        accelerator.save_state(output_dir=os.path.join(self.out_dir_ckpt, 'latest'))
                        unwrapped_model = accelerator.unwrap_model(self.model)
                        if accelerator.is_main_process:
                            accelerator.save_model(unwrapped_model.unet,
                                                   os.path.join(self.out_dir_ckpt, 'latest'), safe_serialization=False)
                            self.save_miscs('latest')

                            # RGB-D joint generation
                            generator = torch.Generator(self.model.device).manual_seed(1024)
                            img = self.model.generate_rgbd(self.prompt, num_inference_steps=50, generator=generator,show_pbar=False)
                            for idx in range(len(self.prompt)):
                                tb_logger.writer.add_image(f'image/{self.prompt[idx]}', img[idx], self.effective_iter)

                            # depth to RGB generation
                            self._depth2image()
                            from diffusers import StableDiffusionControlNetInpaintPipeline
                            # RGB to depth generation
                            self._image2depth()

                        accelerator.wait_for_everyone()

                        accelerator.wait_for_everyone()

                    if self.backup_period > 0 and 0 == self.effective_iter % self.backup_period:
                        unwrapped_model = accelerator.unwrap_model(self.model)
                        if accelerator.is_main_process:
                            unwrapped_model.unet.save_pretrained(
                                os.path.join(self.out_dir_ckpt, self._get_backup_ckpt_name()))
                        accelerator.wait_for_everyone()

                    if self.val_period > 0 and 0 == self.effective_iter % self.val_period:
                        self.validate()

                    # End of training
                    if self.max_iter > 0 and self.effective_iter >= self.max_iter:
                        unwrapped_model = accelerator.unwrap_model(self.model)
                        if accelerator.is_main_process:
                            unwrapped_model.unet.save_pretrained(
                                os.path.join(self.out_dir_ckpt, self._get_backup_ckpt_name()))
                        accelerator.wait_for_everyone()
                        return

                    torch.cuda.empty_cache()
                    # <<< Effective batch end <<<

                    # Epoch end
                    self.n_batch_in_epoch = 0

    def _image2depth(self):
        generator = torch.Generator(self.model.device).manual_seed(1024)
        image2dept_paths = ['/home/aiops/wangzh/data/scannet/scene0593_00/color/000100.jpg',
                            '/home/aiops/wangzh/data/scannet/scene0593_00/color/000700.jpg',
                            '/home/aiops/wangzh/data/scannet/scene0591_01/color/000600.jpg',
                            '/home/aiops/wangzh/data/scannet/scene0591_01/color/001500.jpg']
        for img_idx, image_path in enumerate(image2dept_paths):
            rgb_input = Image.open(image_path)
            depth_pred: MarigoldDepthOutput = self.model.image2depth(
                rgb_input,
                denoising_steps=self.cfg.validation.denoising_steps,
                ensemble_size=self.cfg.validation.ensemble_size,
                processing_res=self.cfg.validation.processing_res,
                match_input_res=self.cfg.validation.match_input_res,
                generator=generator,
                batch_size=self.cfg.validation.ensemble_size,
                # use batch size 1 to increase reproducibility
                color_map="Spectral",
                show_progress_bar=False,
                resample_method=self.cfg.validation.resample_method,
            )
            img = self.model.post_process_rgbd(['None'], [rgb_input], [depth_pred['depth_colored']])
            tb_logger.writer.add_image(f'image2depth_{img_idx}', img[0], self.effective_iter)

    def _depth2image(self):
        generator = torch.Generator(self.model.device).manual_seed(1024)
        if "least_square_disparity" == self.cfg.eval.alignment:
            depth2image_path = ['/home/aiops/wangzh/data/ori_depth_part0-0/sa_10000335.jpg',
                                '/home/aiops/wangzh/data/ori_depth_part0-0/sa_3572319.jpg',
                                '/home/aiops/wangzh/data/ori_depth_part0-0/sa_457934.jpg']
        else:
            depth2image_path = ['/home/aiops/wangzh/data/depth_part0-0/sa_10000335.jpg',
                                '/home/aiops/wangzh/data/depth_part0-0/sa_3572319.jpg',
                                '/home/aiops/wangzh/data/depth_part0-0/sa_457934.jpg']
        prompts = ['Red car parked in the factory',
                  'White gothic church with cemetery next to it',
                  'House with red roof and starry sky in the background']
        for img_idx, depth_path in enumerate(depth2image_path):
            depth_input = Image.open(depth_path)
            image_pred = self.model.single_depth2image(
                depth_input,
                prompts[img_idx],
                num_inference_steps=50,
                processing_res=1024,
                generator=generator,
                show_pbar=False,
                resample_method=self.cfg.validation.resample_method,
            )
            img = self.model.post_process_rgbd([prompts[img_idx]], [image_pred], [depth_input])
            tb_logger.writer.add_image(f'depth2image_{img_idx}', img[0], self.effective_iter)

    def encode_depth(self, depth_in):
        # stack depth into 3-channel
        stacked = self.stack_depth_images(depth_in)
        # encode using VAE encoder
        depth_latent = self.model.encode_rgb(stacked)
        return depth_latent

    @staticmethod
    def stack_depth_images(depth_in):
        if 4 == len(depth_in.shape):
            stacked = depth_in.repeat(1, 3, 1, 1)
        elif 3 == len(depth_in.shape):
            stacked = depth_in.unsqueeze(1)
            stacked = depth_in.repeat(1, 3, 1, 1)
        return stacked

    def _train_step_callback(self):
        """Executed after every iteration"""
        # Save backup (with a larger interval, without training states)
        if self.backup_period > 0 and 0 == self.effective_iter % self.backup_period:
            self.save_checkpoint(
                ckpt_name=self._get_backup_ckpt_name(), save_train_state=False
            )

        _is_latest_saved = False
        # Validation
        if self.val_period > 0 and 0 == self.effective_iter % self.val_period:
            self.in_evaluation = True  # flag to do evaluation in resume run if validation is not finished
            self.save_checkpoint(ckpt_name="latest", save_train_state=True)
            _is_latest_saved = True
            self.validate()
            self.in_evaluation = False
            self.save_checkpoint(ckpt_name="latest", save_train_state=True)

        # Save training checkpoint (can be resumed)
        if (
            self.save_period > 0
            and 0 == self.effective_iter % self.save_period
            and not _is_latest_saved
        ):
            generator = torch.Generator(self.model.device).manual_seed(1024)
            img = self.model.generate_rgbd(self.prompt, num_inference_steps=50, generator=generator, show_pbar=True)
            for idx in range(len(self.prompt)):
                tb_logger.writer.add_image(f'image/{self.prompt[idx]}', img[idx], self.effective_iter)

            self.save_checkpoint(ckpt_name="latest", save_train_state=True)

        # Visualization
        if self.vis_period > 0 and 0 == self.effective_iter % self.vis_period:
            self.visualize()

    def validate(self):
        for i, val_loader in enumerate(self.val_loaders):
            val_dataset_name = val_loader.dataset.disp_name
            val_metric_dic = self.validate_single_dataset(
                data_loader=val_loader, metric_tracker=self.val_metrics
            )

            if self.accelerator.is_main_process:
                val_metric_dic = {k:torch.tensor(v).cuda() for k,v in val_metric_dic.items()}

                tb_logger.log_dic(
                    {f"val/{val_dataset_name}/{k}": v for k, v in val_metric_dic.items()},
                    global_step=self.effective_iter,
                )
                # save to file
                eval_text = eval_dic_to_text(
                    val_metrics=val_metric_dic,
                    dataset_name=val_dataset_name,
                    sample_list_path=val_loader.dataset.filename_ls_path,
                )
                _save_to = os.path.join(
                    self.out_dir_eval,
                    f"eval-{val_dataset_name}-iter{self.effective_iter:06d}.txt",
                )
                with open(_save_to, "w+") as f:
                    f.write(eval_text)

                # Update main eval metric
                if 0 == i:
                    main_eval_metric = val_metric_dic[self.main_val_metric]
                    if (
                        "minimize" == self.main_val_metric_goal
                        and main_eval_metric < self.best_metric
                        or "maximize" == self.main_val_metric_goal
                        and main_eval_metric > self.best_metric
                    ):
                        self.best_metric = main_eval_metric
                        logging.info(
                            f"Best metric: {self.main_val_metric} = {self.best_metric} at iteration {self.effective_iter}"
                        )
                        # Save a checkpoint
                        self.save_checkpoint(
                            ckpt_name='best', save_train_state=False
                        )

            self.accelerator.wait_for_everyone()

    def visualize(self):
        for val_loader in self.vis_loaders:
            vis_dataset_name = val_loader.dataset.disp_name
            vis_out_dir = os.path.join(
                self.out_dir_vis, self._get_backup_ckpt_name(), vis_dataset_name
            )
            os.makedirs(vis_out_dir, exist_ok=True)
            _ = self.validate_single_dataset(
                data_loader=val_loader,
                metric_tracker=self.val_metrics,
                save_to_dir=vis_out_dir,
            )

    @torch.no_grad()
    def validate_single_dataset(
        self,
        data_loader: DataLoader,
        metric_tracker: MetricTracker,
        save_to_dir: str = None,
    ):
        self.model.to(self.device)
        metric_tracker.reset()

        # Generate seed sequence for consistent evaluation
        val_init_seed = self.cfg.validation.init_seed
        val_seed_ls = generate_seed_sequence(val_init_seed, len(data_loader))

        for i, batch in enumerate(
            tqdm(data_loader, desc=f"evaluating on {data_loader.dataset.disp_name}"),
            start=1,
        ):

            rgb_int = batch["rgb_int"]  # [3, H, W]
            # GT depth
            depth_raw_ts = batch["depth_raw_linear"].squeeze()
            depth_raw = depth_raw_ts.cpu().numpy()
            depth_raw_ts = depth_raw_ts.to(self.device)
            valid_mask_ts = batch["valid_mask_raw"].squeeze()
            valid_mask = valid_mask_ts.cpu().numpy()
            valid_mask_ts = valid_mask_ts.to(self.device)

            # Random number generator
            seed = val_seed_ls.pop()
            if seed is None:
                generator = None
            else:
                generator = torch.Generator(device=self.device)
                generator.manual_seed(seed)

            # Predict depth
            pipe_out: MarigoldDepthOutput = self.model.image2depth(
                rgb_int,
                denoising_steps=self.cfg.validation.denoising_steps,
                ensemble_size=self.cfg.validation.ensemble_size,
                processing_res=self.cfg.validation.processing_res,
                match_input_res=self.cfg.validation.match_input_res,
                generator=generator,
                batch_size=self.cfg.validation.ensemble_size,  # use batch size 1 to increase reproducibility
                color_map=None,
                show_progress_bar=False,
                resample_method=self.cfg.validation.resample_method,
            )

            depth_pred: np.ndarray = pipe_out.depth_np

            if "least_square" == self.cfg.eval.alignment:
                depth_pred, scale, shift = align_depth_least_square(
                    gt_arr=depth_raw,
                    pred_arr=depth_pred,
                    valid_mask_arr=valid_mask,
                    return_scale_shift=True,
                    max_resolution=self.cfg.eval.align_max_res,
                )
            else:
                raise RuntimeError(f"Unknown alignment type: {self.cfg.eval.alignment}")

            # Clip to dataset min max
            depth_pred = np.clip(
                depth_pred,
                a_min=data_loader.dataset.min_depth,
                a_max=data_loader.dataset.max_depth,
            )

            # clip to d > 0 for evaluation
            depth_pred = np.clip(depth_pred, a_min=1e-6, a_max=None)

            # Evaluate
            sample_metric = []
            depth_pred_ts = torch.from_numpy(depth_pred).to(self.device)

            for met_func in self.metric_funcs:
                _metric_name = met_func.__name__
                _metric = met_func(depth_pred_ts, depth_raw_ts, valid_mask_ts).cuda(self.accelerator.process_index)
                self.accelerator.wait_for_everyone()
                _metric = self.accelerator.gather_for_metrics(_metric.unsqueeze(0)).mean().item()
                sample_metric.append(_metric.__str__())
                metric_tracker.update(_metric_name, _metric)

            self.accelerator.wait_for_everyone()
            # Save as 16-bit uint png
            if save_to_dir is not None:
                img_name = batch["rgb_relative_path"][0].replace("/", "_")
                png_save_path = os.path.join(save_to_dir, f"{img_name}.png")
                depth_to_save = (pipe_out.depth_np * 65535.0).astype(np.uint16)
                Image.fromarray(depth_to_save).save(png_save_path, mode="I;16")

        return metric_tracker.result()

    def _get_next_seed(self):
        if 0 == len(self.global_seed_sequence):
            self.global_seed_sequence = generate_seed_sequence(
                initial_seed=self.seed,
                length=self.max_iter * self.gradient_accumulation_steps,
            )
            logging.info(
                f"Global seed sequence is generated, length={len(self.global_seed_sequence)}"
            )
        return self.global_seed_sequence.pop()

    def save_miscs(self, ckpt_name):
        ckpt_dir = os.path.join(self.out_dir_ckpt, ckpt_name)
        state = {
            "config": self.cfg,
            "effective_iter": self.effective_iter,
            "epoch": self.epoch,
            "n_batch_in_epoch": self.n_batch_in_epoch,
            "best_metric": self.best_metric,
            "in_evaluation": self.in_evaluation,
            "global_seed_sequence": self.global_seed_sequence,
        }
        train_state_path = os.path.join(ckpt_dir, "trainer.ckpt")
        torch.save(state, train_state_path)

        logging.info(f"Misc state is saved to: {train_state_path}")

    def load_miscs(self, ckpt_path):
        checkpoint = torch.load(os.path.join(ckpt_path, "trainer.ckpt"))
        self.effective_iter = checkpoint["effective_iter"]
        self.epoch = checkpoint["epoch"]
        self.n_batch_in_epoch = checkpoint["n_batch_in_epoch"]
        self.in_evaluation = checkpoint["in_evaluation"]
        self.global_seed_sequence = checkpoint["global_seed_sequence"]

        self.best_metric = checkpoint["best_metric"]

        logging.info(f"Misc state is loaded from {ckpt_path}")


    def save_checkpoint(self, ckpt_name, save_train_state):
        ckpt_dir = os.path.join(self.out_dir_ckpt, ckpt_name)
        logging.info(f"Saving checkpoint to: {ckpt_dir}")
        # Backup previous checkpoint
        temp_ckpt_dir = None
        if os.path.exists(ckpt_dir) and os.path.isdir(ckpt_dir):
            temp_ckpt_dir = os.path.join(
                os.path.dirname(ckpt_dir), f"_old_{os.path.basename(ckpt_dir)}"
            )
            if os.path.exists(temp_ckpt_dir):
                shutil.rmtree(temp_ckpt_dir, ignore_errors=True)
            os.rename(ckpt_dir, temp_ckpt_dir)
            logging.debug(f"Old checkpoint is backed up at: {temp_ckpt_dir}")

        # Save UNet
        unet_path = os.path.join(ckpt_dir, "unet")
        self.model.unet.save_pretrained(unet_path, safe_serialization=False)
        logging.info(f"UNet is saved to: {unet_path}")

        if save_train_state:
            state = {
                "config": self.cfg,
                "effective_iter": self.effective_iter,
                "epoch": self.epoch,
                "n_batch_in_epoch": self.n_batch_in_epoch,
                "best_metric": self.best_metric,
                "in_evaluation": self.in_evaluation,
                "global_seed_sequence": self.global_seed_sequence,
            }
            train_state_path = os.path.join(ckpt_dir, "trainer.ckpt")
            torch.save(state, train_state_path)
            # iteration indicator
            f = open(os.path.join(ckpt_dir, self._get_backup_ckpt_name()), "w")
            f.close()

            logging.info(f"Trainer state is saved to: {train_state_path}")

        # Remove temp ckpt
        if temp_ckpt_dir is not None and os.path.exists(temp_ckpt_dir):
            shutil.rmtree(temp_ckpt_dir, ignore_errors=True)
            logging.debug("Old checkpoint backup is removed.")

    def load_checkpoint(
        self, ckpt_path, load_trainer_state=True, resume_lr_scheduler=True
    ):
        logging.info(f"Loading checkpoint from: {ckpt_path}")
        # Load UNet
        _model_path = os.path.join(ckpt_path, "unet", "diffusion_pytorch_model.bin")
        self.model.unet.load_state_dict(
            torch.load(_model_path, map_location=self.device)
        )
        self.model.unet.to(self.device)
        logging.info(f"UNet parameters are loaded from {_model_path}")

        # Load training states
        if load_trainer_state:
            checkpoint = torch.load(os.path.join(ckpt_path, "trainer.ckpt"))
            self.effective_iter = checkpoint["effective_iter"]
            self.epoch = checkpoint["epoch"]
            self.n_batch_in_epoch = checkpoint["n_batch_in_epoch"]
            self.in_evaluation = checkpoint["in_evaluation"]
            self.global_seed_sequence = checkpoint["global_seed_sequence"]

            self.best_metric = checkpoint["best_metric"]

            self.optimizer.load_state_dict(checkpoint["optimizer"])
            logging.info(f"optimizer state is loaded from {ckpt_path}")

            if resume_lr_scheduler:
                self.lr_scheduler.load_state_dict(checkpoint["lr_scheduler"])
                logging.info(f"LR scheduler state is loaded from {ckpt_path}")

        logging.info(
            f"Checkpoint loaded from: {ckpt_path}. Resume from iteration {self.effective_iter} (epoch {self.epoch})"
        )
        return

    def _get_backup_ckpt_name(self):
        return f"iter_{self.effective_iter:06d}"