Upload train_stage2.py with huggingface_hub

train_stage2.py

@@ -0,0 +1,290 @@
+"""
+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)
+    
+