""" training script for imagedream - the config system is similar with stable diffusion ldm code base(using omigaconf, yaml; target, params initialization, etc.) - the training code base is similar with unidiffuser training code base using accelerate concat channel as input, pred xyz value mapped pixedl as groundtruth """ from omegaconf import OmegaConf import argparse import datetime from pathlib import Path from torch.utils.data import DataLoader import os.path as osp import numpy as np import os import torch import wandb from libs.base_utils import get_data_generator, PrintContext from libs.base_utils import setup, instantiate_from_config, dct2str, add_prefix, get_obj_from_str from absl import logging from einops import rearrange from libs.sample import ImageDreamDiffusion def train(config, unk): # using pipeline to extract models accelerator, device = setup(config, unk) with PrintContext(f"{'access STAT':-^50}", accelerator.is_main_process): print(accelerator.state) dtype = { "fp16": torch.float16, "fp32": torch.float32, "no": torch.float32, "bf16": torch.bfloat16, }[accelerator.state.mixed_precision] num_frames = config.num_frames ################## load models ################## model_config = config.models.config model_config = OmegaConf.load(model_config) model = instantiate_from_config(model_config.model) state_dict = torch.load(config.models.resume, map_location="cpu") model_in_conv_keys = ["model.diffusion_model.input_blocks.0.0.weight",] in_conv_keys = ["diffusion_model.input_blocks.0.0.weight"] def modify_keys(state_dict, in_keys, out_keys, cur_state_dict=None): print("this function only for fuse channel model") for in_key in in_keys: p = state_dict[in_key] if cur_state_dict is not None: p_cur = cur_state_dict[in_key] print(p_cur.shape, p.shape) if p_cur.shape == p.shape: print(f"skip {in_key} because of same shape") continue state_dict[in_key] = torch.cat([p, torch.zeros_like(p)], dim=1) * 0.5 for out_key in out_keys: p = state_dict[out_key] if cur_state_dict is not None: p_cur = cur_state_dict[out_key] print(p_cur.shape, p.shape) if p_cur.shape == p.shape: print(f"skip {out_key} because of same shape") continue state_dict[out_key] = torch.cat([p, torch.zeros_like(p)], dim=0) return state_dict def wipe_keys(state_dict, keys): for key in keys: state_dict.pop(key) return state_dict unet_config = model_config.model.params.unet_config is_normal_inout_channel = not (unet_config.params.in_channels != 4 or unet_config.params.out_channels != 4) if not is_normal_inout_channel: state_dict = modify_keys(state_dict, model_in_conv_keys, [], model.state_dict()) print(model.load_state_dict(state_dict, strict=False)) print("loaded model from {}".format(config.models.resume)) if config.models.get("resume_unet", None) is not None: unet_state_dict = torch.load(config.models.resume_unet, map_location="cpu") if not is_normal_inout_channel: unet_state_dict = modify_keys(unet_state_dict, in_conv_keys, [], model.model.state_dict()) print(model.model.load_state_dict(unet_state_dict, strict= False)) print(f"______ load unet from {config.models.resume_unet} ______") model.to(device) model.device = device model.clip_model.device = device ################# setup optimizer ################# from torch.optim import AdamW from accelerate.utils import DummyOptim optimizer_cls = ( AdamW if accelerator.state.deepspeed_plugin is None or "optimizer" not in accelerator.state.deepspeed_plugin.deepspeed_config else DummyOptim ) optimizer = optimizer_cls(model.model.parameters(), **config.optimizer) ################# prepare datasets ################# dataset = instantiate_from_config(config.train_data) eval_dataset = instantiate_from_config(config.eval_data) dl_config = config.dataloader dataloader = DataLoader(dataset, **dl_config, batch_size=config.batch_size) model, optimizer, dataloader, = accelerator.prepare(model, optimizer, dataloader) generator = get_data_generator(dataloader, accelerator.is_main_process, "train") if config.get("sampler", None) is not None: sampler_cls = get_obj_from_str(config.sampler.target) sampler = sampler_cls(model, device, dtype, **config.sampler.params) else: sampler = ImageDreamDiffusion(model, config.mode, num_frames, device, dtype, dataset.camera_views, offset_noise=config.get("offset_noise", False), ref_position=dataset.ref_position, random_background=dataset.random_background, resize_rate=dataset.resize_rate) ################# evaluation code ################# def evaluation(): from PIL import Image import numpy as np return_ls = [] for i in range(accelerator.process_index, len(eval_dataset), accelerator.num_processes): item = eval_dataset[i] cond = item['cond'] images = sampler.diffuse("3D assets.", cond, pixel_images=item["cond_raw_images"], n_test=2) images = np.concatenate(images, 0) images = [Image.fromarray(images)] return_ls.append(dict(images=images, ident=eval_dataset[i]['ident'])) return return_ls global_step = 0 total_step = 0 log_step = 0 eval_step = 0 save_step = config.save_interval unet = model.model while True: item = next(generator) unet.train() bs = item["clip_cond"].shape[0] BS = bs * num_frames item["clip_cond"] = item["clip_cond"].to(device).to(dtype) item["vae_cond"] = item["vae_cond"].to(device).to(dtype) camera_input = item["cameras"].to(device) camera_input = camera_input.reshape((BS, camera_input.shape[-1])) gd_type = config.get("gd_type", "pixel") if gd_type == "pixel": item["target_images_vae"] = item["target_images_vae"].to(device).to(dtype) gd = item["target_images_vae"] elif gd_type == "xyz": item["target_images_xyz_vae"] = item["target_images_xyz_vae"].to(device).to(dtype) item["target_images_vae"] = item["target_images_vae"].to(device).to(dtype) gd = item["target_images_xyz_vae"] elif gd_type == "fusechannel": item["target_images_vae"] = item["target_images_vae"].to(device).to(dtype) item["target_images_xyz_vae"] = item["target_images_xyz_vae"].to(device).to(dtype) gd = torch.cat((item["target_images_vae"], item["target_images_xyz_vae"]), dim=0) else: raise NotImplementedError with torch.no_grad(), accelerator.autocast("cuda"): ip_embed = model.clip_model.encode_image_with_transformer(item["clip_cond"]) ip_ = ip_embed.repeat_interleave(num_frames, dim=0) ip_img = model.get_first_stage_encoding(model.encode_first_stage(item["vae_cond"])) gd = rearrange(gd, "B F C H W -> (B F) C H W") pixel_images = rearrange(item["target_images_vae"], "B F C H W -> (B F) C H W") latent_target_images = model.get_first_stage_encoding(model.encode_first_stage(gd)) pixel_images = model.get_first_stage_encoding(model.encode_first_stage(pixel_images)) if gd_type == "fusechannel": latent_target_images = rearrange(latent_target_images, "(B F) C H W -> B F C H W", B=bs * 2) image_latent, xyz_latent = torch.chunk(latent_target_images, 2) fused_channel_latent = torch.cat((image_latent, xyz_latent), dim=-3) latent_target_images = rearrange(fused_channel_latent, "B F C H W -> (B F) C H W") if item.get("captions", None) is not None: caption_ls = np.array(item["caption"]).T.reshape((-1, BS)).squeeze() prompt_cond = model.get_learned_conditioning(caption_ls) elif item.get("caption", None) is not None: prompt_cond = model.get_learned_conditioning(item["caption"]) prompt_cond = prompt_cond.repeat_interleave(num_frames, dim=0) else: prompt_cond = model.get_learned_conditioning(["3D assets."]).repeat(BS, 1, 1) condition = { "context": prompt_cond, "ip": ip_, # "ip_img": ip_img, "camera": camera_input, "pixel_images": pixel_images, } with torch.autocast("cuda"), accelerator.accumulate(model): time_steps = torch.randint(0, model.num_timesteps, (BS,), device=device) noise = torch.randn_like(latent_target_images, device=device) x_noisy = model.q_sample(latent_target_images, time_steps, noise) output = unet(x_noisy, time_steps, **condition, num_frames=num_frames) loss = torch.nn.functional.mse_loss(noise, output) accelerator.backward(loss) optimizer.step() optimizer.zero_grad() global_step += 1 total_step = global_step * config.total_batch_size if total_step > log_step: metrics = dict( loss = accelerator.gather(loss.detach().mean()).mean().item(), scale = accelerator.scaler.get_scale() if accelerator.scaler is not None else -1 ) log_step += config.log_interval if accelerator.is_main_process: logging.info(dct2str(dict(step=total_step, **metrics))) wandb.log(add_prefix(metrics, 'train'), step=total_step) if total_step > save_step and accelerator.is_main_process: logging.info("saving done") torch.save(unet.state_dict(), osp.join(config.ckpt_root, f"unet-{total_step}")) save_step += config.save_interval logging.info("save done") if total_step > eval_step: logging.info("evaluationing") unet.eval() return_ls = evaluation() cur_eval_base = osp.join(config.eval_root, f"{total_step:07d}") os.makedirs(cur_eval_base, exist_ok=True) wandb_image_ls = [] for item in return_ls: for i, im in enumerate(item["images"]): im.save(osp.join(cur_eval_base, f"{item['ident']}-{i:03d}-{accelerator.process_index}-.png")) wandb_image_ls.append(wandb.Image(im, caption=f"{item['ident']}-{i:03d}-{accelerator.process_index}")) wandb.log({"eval_samples": wandb_image_ls}) eval_step += config.eval_interval logging.info("evaluation done") accelerator.wait_for_everyone() if total_step > config.max_step: break if __name__ == "__main__": # load config from config path, then merge with cli args parser = argparse.ArgumentParser() parser.add_argument( "--config", type=str, default="configs/nf7_v3_SNR_rd_size_stroke.yaml" ) parser.add_argument( "--logdir", type=str, default="train_logs", help="the dir to put logs" ) parser.add_argument( "--resume_workdir", type=str, default=None, help="specify to do resume" ) args, unk = parser.parse_known_args() print(args, unk) config = OmegaConf.load(args.config) if args.resume_workdir is not None: assert osp.exists(args.resume_workdir), f"{args.resume_workdir} not exists" config.config.workdir = args.resume_workdir config.config.resume = True OmegaConf.set_struct(config, True) # prevent adding new keys cli_conf = OmegaConf.from_cli(unk) config = OmegaConf.merge(config, cli_conf) config = config.config OmegaConf.set_struct(config, False) config.logdir = args.logdir config.config_name = Path(args.config).stem train(config, unk)