content: # Hunyuan 3D is licensed under the TENCENT HUNYUAN NON-COMMERCIAL LICENSE AGREEMENT
  # except for the third-party components listed below.
  # Hunyuan 3D does not impose any additional limitations beyond what is outlined
  # in the repsective licenses of these third-party components.
  # Users must comply with all terms and conditions of original licenses of these third-party
  # components and must ensure that the usage of the third party components adheres to
  # all relevant laws and regulations.
  
  # For avoidance of doubts, Hunyuan 3D means the large language models and
  # their software and algorithms, including trained model weights, parameters (including
  # optimizer states), machine-learning model code, inference-enabling code, training-enabling code,
  # fine-tuning enabling code and other elements of the foregoing made publicly available
  # by Tencent in accordance with TENCENT HUNYUAN COMMUNITY LICENSE AGREEMENT.
  
  import torch
  import os, sys
  import argparse
  import shutil
  import subprocess
  from omegaconf import OmegaConf
  
  from pytorch_lightning import seed_everything
  from pytorch_lightning.trainer import Trainer
  from pytorch_lightning.strategies import DDPStrategy
  from pytorch_lightning.callbacks import Callback
  from pytorch_lightning.utilities import rank_zero_only, rank_zero_warn
  
  from src.utils.train_util import instantiate_from_config
  import warnings
  
  warnings.filterwarnings("ignore")
  from diffusers.utils import logging as diffusers_logging
  
  diffusers_logging.set_verbosity(50)
  
  
  @rank_zero_only
  def rank_zero_print(*args):
      print(*args)
  
  
  def get_parser(**parser_kwargs):
      def str2bool(v):
          if isinstance(v, bool):
              return v
          if v.lower() in ("yes", "true", "t", "y", "1"):
              return True
          elif v.lower() in ("no", "false", "f", "n", "0"):
              return False
          else:
              raise argparse.ArgumentTypeError("Boolean value expected.")
  
      parser = argparse.ArgumentParser(**parser_kwargs)
      parser.add_argument(
          "-r",
          "--resume",
          type=str,
          default=None,
          help="resume from checkpoint",
      )
      parser.add_argument(
          "--resume_weights_only",
          action="store_true",
          help="only resume model weights",
      )
      parser.add_argument(
          "-b",
          "--base",
          type=str,
          default="base_config.yaml",
          help="path to base configs",
      )
      parser.add_argument(
          "-n",
          "--name",
          type=str,
          default="",
          help="experiment name",
      )
      parser.add_argument(
          "--num_nodes",
          type=int,
          default=1,
          help="number of nodes to use",
      )
      parser.add_argument(
          "--gpus",
          type=str,
          default="0,",
          help="gpu ids to use",
      )
      parser.add_argument(
          "-s",
          "--seed",
          type=int,
          default=42,
          help="seed for seed_everything",
      )
      parser.add_argument(
          "-l",
          "--logdir",
          type=str,
          default="logs",
          help="directory for logging data",
      )
      return parser
  
  
  class SetupCallback(Callback):
      def __init__(self, resume, logdir, ckptdir, cfgdir, config):
          super().__init__()
          self.resume = resume
          self.logdir = logdir
          self.ckptdir = ckptdir
          self.cfgdir = cfgdir
          self.config = config
  
      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)
  
              rank_zero_print("Project config")
              rank_zero_print(OmegaConf.to_yaml(self.config))
              OmegaConf.save(self.config, os.path.join(self.cfgdir, "project.yaml"))
  
  
  class CodeSnapshot(Callback):
      """
      Modified from https://github.com/threestudio-project/threestudio/blob/main/threestudio/utils/callbacks.py#L60
      """
  
      def __init__(self, savedir):
          self.savedir = savedir
  
      def get_file_list(self):
          return [
              b.decode()
              for b in set(subprocess.check_output('git ls-files -- ":!:configs/*"', shell=True).splitlines())
              | set(  # hard code, TODO: use config to exclude folders or files
                  subprocess.check_output("git ls-files --others --exclude-standard", shell=True).splitlines()
              )
          ]
  
      @rank_zero_only
      def save_code_snapshot(self):
          os.makedirs(self.savedir, exist_ok=True)
  
      #       for f in self.get_file_list():
      #           if not os.path.exists(f) or os.path.isdir(f):
      #               continue
      #           os.makedirs(os.path.join(self.savedir, os.path.dirname(f)), exist_ok=True)
      #           shutil.copyfile(f, os.path.join(self.savedir, f))
  
      def on_fit_start(self, trainer, pl_module):
          try:
              self.save_code_snapshot()
          except:
              rank_zero_warn(
                  "Code snapshot is not saved. Please make sure you have git installed and are in a git repository."
              )
  
  
  if __name__ == "__main__":
      # add cwd for convenience and to make classes in this file available when
      # running as `python main.py`
      sys.path.append(os.getcwd())
      torch.set_float32_matmul_precision("medium")
  
      parser = get_parser()
      opt, unknown = parser.parse_known_args()
  
      cfg_fname = os.path.split(opt.base)[-1]
      cfg_name = os.path.splitext(cfg_fname)[0]
      exp_name = "-" + opt.name if opt.name != "" else ""
      logdir = os.path.join(opt.logdir, cfg_name + exp_name)
  
      # assert not os.path.exists(logdir) or 'test' in logdir, logdir
      if os.path.exists(logdir) and opt.resume is None:
          auto_resume_path = os.path.join(logdir, "checkpoints", "last.ckpt")
          if os.path.exists(auto_resume_path):
              opt.resume = auto_resume_path
              print(f"Auto set resume ckpt {opt.resume}")
  
      ckptdir = os.path.join(logdir, "checkpoints")
      cfgdir = os.path.join(logdir, "configs")
      codedir = os.path.join(logdir, "code")
  
      node_rank = int(os.environ.get("NODE_RANK", 0))  # 当前节点的编号
      local_rank = int(os.environ.get("LOCAL_RANK", 0))  # 当前节点上的 GPU 编号
      num_gpus_per_node = torch.cuda.device_count()  # 每个节点上的 GPU 数量
  
      global_rank = node_rank * num_gpus_per_node + local_rank
      seed_everything(opt.seed + global_rank)
  
      # init configs
      config = OmegaConf.load(opt.base)
      lightning_config = config.lightning
      trainer_config = lightning_config.trainer
  
      trainer_config["accelerator"] = "gpu"
      rank_zero_print(f"Running on GPUs {opt.gpus}")
      try:
          ngpu = int(opt.gpus)
      except:
          ngpu = len(opt.gpus.strip(",").split(","))
      trainer_config["devices"] = ngpu
  
      trainer_opt = argparse.Namespace(**trainer_config)
      lightning_config.trainer = trainer_config
  
      # model
      model = instantiate_from_config(config.model)
  
      model_unet = model.unet.unet
      model_unet_prefix = "unet.unet."
      if hasattr(model_unet, "unet"):
          model_unet = model_unet.unet
          model_unet_prefix += "unet."
  
      if getattr(config, "init_unet_from", None):
          unet_ckpt_path = config.init_unet_from
          sd = torch.load(unet_ckpt_path, map_location="cpu")
          model_unet.load_state_dict(sd, strict=True)
  
      if getattr(config, "init_vae_from", None):
          vae_ckpt_path = config.init_vae_from
          sd_vae = torch.load(vae_ckpt_path, map_location="cpu")
  
          def replace_key(key_str):
              replace_pairs = [("key", "to_k"), ("query", "to_q"), ("value", "to_v"), ("proj_attn", "to_out.0")]
              for replace_pair in replace_pairs:
                  key_str = key_str.replace(replace_pair[0], replace_pair[1])
              return key_str
  
          sd_vae = {replace_key(k): v for k, v in sd_vae.items()}
          model.pipeline.vae.load_state_dict(sd_vae, strict=True)
  
      if hasattr(model.unet, "controlnet"):
          if getattr(config, "init_control_from", None):
              unet_ckpt_path = config.init_control_from
              sd_control = torch.load(unet_ckpt_path, map_location="cpu")
              model.unet.controlnet.load(sd_control, strict=True)
  
      noise_in_channels = config.model.params.get("noise_in_channels", None)
      if noise_in_channels is not None:
          with torch.no_grad():
              new_conv_in = torch.nn.Conv2d(
                  noise_in_channels,
                  model_unet.conv_in.out_channels,
                  model_unet.conv_in.kernel_size,
                  model_unet.conv_in.stride,
                  model_unet.conv_in.padding,
              )
              new_conv_in.weight.zero_()
              new_conv_in.weight[:, : model_unet.conv_in.in_channels, :, :].copy_(model_unet.conv_in.weight)
  
              new_conv_in.bias.zero_()
              new_conv_in.bias[: model_unet.conv_in.bias.size(0)].copy_(model_unet.conv_in.bias)
  
              model_unet.conv_in = new_conv_in
  
      if hasattr(model.unet, "controlnet"):
          if config.model.params.get("control_in_channels", None):
              control_in_channels = config.model.params.control_in_channels
              model.unet.controlnet.config["conditioning_channels"] = control_in_channels
              condition_conv_in = model.unet.controlnet.controlnet_cond_embedding.conv_in
  
              new_condition_conv_in = torch.nn.Conv2d(
                  control_in_channels,
                  condition_conv_in.out_channels,
                  kernel_size=condition_conv_in.kernel_size,
                  stride=condition_conv_in.stride,
                  padding=condition_conv_in.padding,
              )
  
              with torch.no_grad():
                  new_condition_conv_in.weight[:, : condition_conv_in.in_channels, :, :] = condition_conv_in.weight
                  if condition_conv_in.bias is not None:
                      new_condition_conv_in.bias = condition_conv_in.bias
  
              model.unet.controlnet.controlnet_cond_embedding.conv_in = new_condition_conv_in
  
      rank_zero_print(f"Loaded Init ...")
  
      if getattr(config, "resume_from", None):
          cnet_ckpt_path = config.resume_from
          sds = torch.load(cnet_ckpt_path, map_location="cpu")["state_dict"]
          sd0 = {k[len(model_unet_prefix) :]: v for k, v in sds.items() if model_unet_prefix in k}
          # model.unet.unet.unet.load_state_dict(sd0, strict=True)
          model_unet.load_state_dict(sd0, strict=True)
          if hasattr(model.unet, "controlnet"):
              sd1 = {k[16:]: v for k, v in sds.items() if "unet.controlnet." in k}
              model.unet.controlnet.load_state_dict(sd1, strict=True)
          rank_zero_print(f"Loaded {cnet_ckpt_path} ...")
  
      if opt.resume and opt.resume_weights_only:
          model = model.__class__.load_from_checkpoint(opt.resume, **config.model.params)
  
      model.logdir = logdir
  
      # trainer and callbacks
      trainer_kwargs = dict()
  
      # logger
      default_logger_cfg = {
          "target": "pytorch_lightning.loggers.TensorBoardLogger",
          "params": {
              "name": "tensorboard",
              "save_dir": logdir,
              "version": "0",
          },
      }
      logger_cfg = OmegaConf.merge(default_logger_cfg)
      trainer_kwargs["logger"] = instantiate_from_config(logger_cfg)
  
      # model checkpoint
      default_modelckpt_cfg = {
          "target": "pytorch_lightning.callbacks.ModelCheckpoint",
          "params": {
              "dirpath": ckptdir,
              "filename": "{step:08}",
              "verbose": True,
              "save_last": True,
              "every_n_train_steps": 5000,
              "save_top_k": -1,  # save all checkpoints
          },
      }
  
      if "modelcheckpoint" in lightning_config:
          modelckpt_cfg = lightning_config.modelcheckpoint
      else:
          modelckpt_cfg = OmegaConf.create()
      modelckpt_cfg = OmegaConf.merge(default_modelckpt_cfg, modelckpt_cfg)
  
      # callbacks
      default_callbacks_cfg = {
          "setup_callback": {
              "target": "train.SetupCallback",
              "params": {
                  "resume": opt.resume,
                  "logdir": logdir,
                  "ckptdir": ckptdir,
                  "cfgdir": cfgdir,
                  "config": config,
              },
          },
          "learning_rate_logger": {
              "target": "pytorch_lightning.callbacks.LearningRateMonitor",
              "params": {
                  "logging_interval": "step",
              },
          },
          "code_snapshot": {
              "target": "train.CodeSnapshot",
              "params": {
                  "savedir": codedir,
              },
          },
      }
      default_callbacks_cfg["checkpoint_callback"] = modelckpt_cfg
  
      if "callbacks" in lightning_config:
          callbacks_cfg = lightning_config.callbacks
      else:
          callbacks_cfg = OmegaConf.create()
      callbacks_cfg = OmegaConf.merge(default_callbacks_cfg, callbacks_cfg)
  
      trainer_kwargs["callbacks"] = [instantiate_from_config(callbacks_cfg[k]) for k in callbacks_cfg]
  
      trainer_kwargs["precision"] = "bf16"
      trainer_kwargs["strategy"] = DDPStrategy(find_unused_parameters=False)
  
      # trainer
      trainer = Trainer(**trainer_config, **trainer_kwargs, num_nodes=opt.num_nodes, inference_mode=False)
      trainer.logdir = logdir
  
      # data
      data = instantiate_from_config(config.data)
      data.prepare_data()
      data.setup("fit")
  
      # configure learning rate
      base_lr = config.model.base_learning_rate
      if "accumulate_grad_batches" in lightning_config.trainer:
          accumulate_grad_batches = lightning_config.trainer.accumulate_grad_batches
      else:
          accumulate_grad_batches = 1
      rank_zero_print(f"accumulate_grad_batches = {accumulate_grad_batches}")
      lightning_config.trainer.accumulate_grad_batches = accumulate_grad_batches
      model.learning_rate = base_lr
      rank_zero_print("++++ NOT USING LR SCALING ++++")
      rank_zero_print(f"Setting learning rate to {model.learning_rate:.2e}")
  
      # run training loop
      if opt.resume and not opt.resume_weights_only:
          trainer.fit(model, data, ckpt_path=opt.resume)
      else:
          trainer.fit(model, data)
  