Upload train.py with huggingface_hub

train.py

@@ -0,0 +1,354 @@
+"""
+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
+
+"""
+
+from omegaconf import OmegaConf
+import argparse
+from pathlib import Path
+from torch.utils.data import DataLoader
+import os.path as osp
+import numpy as np
+import os
+import torch
+from PIL import Image
+import numpy as np
+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 imagedream.camera_utils import get_camera
+from libs.sample import ImageDreamDiffusion
+from rich import print
+
+
+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")
+
+    print(model.load_state_dict(state_dict, strict=False))
+    print("loaded model from {}".format(config.models.resume))
+
+    latest_step = 0
+    if config.get("resume", False):
+        print("resuming from specified workdir")
+        ckpts = os.listdir(config.ckpt_root)
+        if len(ckpts) == 0:
+            print("no ckpt found")
+        else:
+            latest_ckpt = sorted(ckpts, key=lambda x: int(x.split("-")[-1]))[-1]
+            latest_step = int(latest_ckpt.split("-")[-1])
+            print("loadding ckpt from ", osp.join(config.ckpt_root, latest_ckpt))
+            unet_state_dict = torch.load(
+                osp.join(config.ckpt_root, latest_ckpt), map_location="cpu"
+            )
+            print(model.model.load_state_dict(unet_state_dict, strict=False))
+
+    elif config.models.get("resume_unet", None) is not None:
+        unet_state_dict = torch.load(config.models.resume_unet, map_location="cpu")
+        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)
+    in_the_wild_images = (
+        instantiate_from_config(config.in_the_wild_images)
+        if config.get("in_the_wild_images", None) is not None
+        else None
+    )
+
+    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,
+            mode=config.mode,
+            num_frames=num_frames,
+            device=device,
+            dtype=dtype,
+            camera_views=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():
+        return_ls = []
+        for i in range(
+            accelerator.process_index, len(eval_dataset), accelerator.num_processes
+        ):
+            cond = eval_dataset[i]["cond"]
+
+            images = sampler.diffuse("3D assets.", cond, 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
+
+    def evaluation2():
+        # eval for common used in the wild image
+        return_ls = []
+        in_the_wild_images.init_item()
+        for i in range(
+            accelerator.process_index,
+            len(in_the_wild_images),
+            accelerator.num_processes,
+        ):
+            cond = in_the_wild_images[i]["cond"]
+            images = sampler.diffuse("3D assets.", cond, n_test=2)
+            images = np.concatenate(images, 0)
+            images = [Image.fromarray(images)]
+            return_ls.append(dict(images=images, ident=in_the_wild_images[i]["ident"]))
+        return return_ls
+
+    if latest_step == 0:
+        global_step = 0
+        total_step = 0
+        log_step = 0
+        eval_step = 0
+        save_step = 0
+    else:
+        global_step = latest_step // config.total_batch_size
+        total_step = latest_step
+        log_step = latest_step + config.log_interval
+        eval_step = latest_step + config.eval_interval
+        save_step = latest_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)
+            )
+            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")
+
+            latent_target_images = model.get_first_stage_encoding(
+                model.encode_first_stage(gd)
+            )
+
+            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,
+            }
+
+        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)
+            # noise_img, _ = torch.chunk(noise, 2, dim=1)
+            # noise = torch.cat((noise_img, noise_img), dim=1)
+            x_noisy = model.q_sample(latent_target_images, time_steps, noise)
+            output = unet(x_noisy, time_steps, **condition, num_frames=num_frames)
+            reshaped_pred = output.reshape(bs, num_frames, *output.shape[1:]).permute(
+                1, 0, 2, 3, 4
+            )
+            reshaped_noise = noise.reshape(bs, num_frames, *noise.shape[1:]).permute(
+                1, 0, 2, 3, 4
+            )
+            true_pred = reshaped_pred[: num_frames - 1]
+            fake_pred = reshaped_pred[num_frames - 1 :]
+            true_noise = reshaped_noise[: num_frames - 1]
+            fake_noise = reshaped_noise[num_frames - 1 :]
+            loss = (
+                torch.nn.functional.mse_loss(true_noise, true_pred)
+                + torch.nn.functional.mse_loss(fake_noise, fake_pred) * 0
+            )
+
+            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)
+            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",
+                        )
+                    )
+
+            return_ls2 = evaluation2()
+            cur_eval_base = osp.join(config.eval_root2, f"{total_step:07d}")
+            os.makedirs(cur_eval_base, exist_ok=True)
+            for item in return_ls2:
+                for i, im in enumerate(item["images"]):
+                    im.save(
+                        osp.join(
+                            cur_eval_base,
+                            f"{item['ident']}-{i:03d}-{accelerator.process_index}-inthewild.png",
+                        )
+                    )
+            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)