"""
Generate a large batch of image samples from a model and save them as a large
numpy array. This can be used to produce samples for FID evaluation.
"""
import argparse
import json
import sys
import os
sys.path.append('.')
from pdb import set_trace as st
import imageio
import numpy as np
import torch as th
import torch.distributed as dist
from guided_diffusion import dist_util, logger
from guided_diffusion.script_util import (
NUM_CLASSES,
model_and_diffusion_defaults,
create_model_and_diffusion,
add_dict_to_argparser,
args_to_dict,
continuous_diffusion_defaults,
control_net_defaults,
)
from pathlib import Path
from tqdm import tqdm, trange
import dnnlib
from dnnlib.util import EasyDict, InfiniteSampler
from nsr.train_util_diffusion import TrainLoop3DDiffusion as TrainLoop
from guided_diffusion.continuous_diffusion import make_diffusion as make_sde_diffusion
import nsr
import nsr.lsgm
from nsr.script_util import create_3DAE_model, encoder_and_nsr_defaults, loss_defaults, AE_with_Diffusion, rendering_options_defaults, eg3d_options_default, dataset_defaults
from datasets.shapenet import load_eval_data
from torch.utils.data import Subset
from datasets.eg3d_dataset import init_dataset_kwargs
from datasets.eg3d_dataset import LMDBDataset_MV_Compressed_eg3d
SEED = 0
def main(args):
# args.rendering_kwargs = rendering_options_defaults(args)
dist_util.setup_dist(args)
logger.configure(dir=args.logdir)
th.cuda.empty_cache()
th.cuda.manual_seed_all(SEED)
np.random.seed(SEED)
# * set denoise model args
logger.log("creating model and diffusion...")
args.img_size = [args.image_size_encoder]
# ! no longer required for LDM
# args.denoise_in_channels = args.out_chans
# args.denoise_out_channels = args.out_chans
args.image_size = args.image_size_encoder # 224, follow the triplane size
denoise_model, diffusion = create_model_and_diffusion(
**args_to_dict(args,
model_and_diffusion_defaults().keys()))
if 'cldm' in args.trainer_name:
assert isinstance(denoise_model, tuple)
denoise_model, controlNet = denoise_model
controlNet.to(dist_util.dev())
controlNet.train()
else:
controlNet = None
opts = eg3d_options_default()
if args.sr_training:
args.sr_kwargs = dnnlib.EasyDict(
channel_base=opts.cbase,
channel_max=opts.cmax,
fused_modconv_default='inference_only',
use_noise=True
) # ! close noise injection? since noise_mode='none' in eg3d
# denoise_model.load_state_dict(
# dist_util.load_state_dict(args.ddpm_model_path, map_location="cpu"))
denoise_model.to(dist_util.dev())
if args.use_fp16:
denoise_model.convert_to_fp16()
denoise_model.eval()
# * auto-encoder reconstruction model
logger.log("creating 3DAE...")
auto_encoder = create_3DAE_model(
**args_to_dict(args,
encoder_and_nsr_defaults().keys()))
# logger.log("AE triplane decoder reuses G_ema decoder...")
# auto_encoder.decoder.register_buffer('w_avg', G_ema.backbone.mapping.w_avg)
# print(auto_encoder.decoder.w_avg.shape) # [512]
# auto_encoder.load_state_dict(
# dist_util.load_state_dict(args.rec_model_path, map_location="cpu"))
auto_encoder.to(dist_util.dev())
auto_encoder.eval()
# TODO, how to set the scale?
logger.log("create dataset")
# data = None
if args.objv_dataset:
from datasets.g_buffer_objaverse import load_data, load_eval_data, load_memory_data, load_wds_data
else: # shapenet
from datasets.shapenet import load_data, load_eval_data, load_memory_data
eval_data = None
# if args.cfg in ('afhq', 'ffhq'):
# # ! load data
# if args.use_lmdb:
# logger.log("creating LMDB eg3d data loader...")
# training_set = LMDBDataset_MV_Compressed_eg3d(
# args.data_dir,
# args.image_size,
# args.image_size_encoder,
# )
# else:
# logger.log("creating eg3d data loader...")
# training_set_kwargs, dataset_name = init_dataset_kwargs(
# data=args.data_dir,
# class_name='datasets.eg3d_dataset.ImageFolderDataset'
# ) # only load pose here
# # if args.cond and not training_set_kwargs.use_labels:
# # raise Exception('check here')
# # training_set_kwargs.use_labels = args.cond
# training_set_kwargs.use_labels = True
# training_set_kwargs.xflip = True
# training_set_kwargs.random_seed = SEED
# # desc = f'{args.cfg:s}-{dataset_name:s}-gpus{c.num_gpus:d}-batch{c.batch_size:d}-gamma{c.loss_kwargs.r1_gamma:g}'
# # * construct ffhq/afhq dataset
# training_set = dnnlib.util.construct_class_by_name(
# **training_set_kwargs) # subclass of training.dataset.Dataset
# training_set = dnnlib.util.construct_class_by_name(
# **training_set_kwargs) # subclass of training.dataset.Dataset
# # training_set_sampler = InfiniteSampler(
# # dataset=training_set,
# # rank=dist_util.get_rank(),
# # num_replicas=dist_util.get_world_size(),
# # seed=SEED)
# # data = iter(
# # th.utils.data.DataLoader(dataset=training_set,
# # sampler=training_set_sampler,
# # batch_size=args.batch_size,
# # pin_memory=True,
# # num_workers=args.num_workers,))
# # # prefetch_factor=2))
# # training_set_sampler = InfiniteSampler(
# # dataset=training_set,
# # rank=dist_util.get_rank(),
# # num_replicas=dist_util.get_world_size(),
# # seed=SEED)
# # data = iter(
# # th.utils.data.DataLoader(
# # dataset=training_set,
# # sampler=training_set_sampler,
# # batch_size=args.batch_size,
# # pin_memory=True,
# # num_workers=args.num_workers,
# # persistent_workers=args.num_workers > 0,
# # # prefetch_factor=max(8//args.batch_size, 2),
# # ))
# eval_data = th.utils.data.DataLoader(dataset=Subset(
# training_set, np.arange(25)),
# batch_size=args.eval_batch_size,
# num_workers=1)
# else:
# logger.log("creating data loader...")
# # if args.objv_dataset:
# # from datasets.g_buffer_objaverse import load_data, load_eval_data, load_memory_data
# # else: # shapenet
# # from datasets.shapenet import load_data, load_eval_data, load_memory_data
# # eval_data = load_eval_data(
# # file_path=args.eval_data_dir,
# # batch_size=args.eval_batch_size,
# # reso=args.image_size,
# # reso_encoder=args.image_size_encoder, # 224 -> 128
# # num_workers=args.num_workers,
# # load_depth=True, # for evaluation
# # interval=args.interval,
# # use_lmdb=args.use_lmdb,
# # )
# if args.use_wds:
# if args.eval_data_dir == 'NONE':
# with open(args.eval_shards_lst) as f:
# eval_shards_lst = [url.strip() for url in f.readlines()]
# else:
# eval_shards_lst = args.eval_data_dir # auto expanded
# eval_data = load_wds_data(
# eval_shards_lst, args.image_size, args.image_size_encoder,
# args.eval_batch_size, args.num_workers,
# **args_to_dict(args,
# dataset_defaults().keys()))
# else:
# eval_data = load_eval_data(
# file_path=args.eval_data_dir,
# batch_size=args.eval_batch_size,
# reso=args.image_size,
# reso_encoder=args.image_size_encoder, # 224 -> 128
# num_workers=args.num_workers,
# # load_depth=True, # for evaluation
# **args_to_dict(args,
# dataset_defaults().keys()))
TrainLoop = {
'adm': nsr.TrainLoop3DDiffusion,
'vpsde_crossattn': nsr.lsgm.TrainLoop3DDiffusionLSGM_crossattn,
}[args.trainer_name]
# continuous
if 'vpsde' in args.trainer_name:
sde_diffusion = make_sde_diffusion(
dnnlib.EasyDict(
args_to_dict(args,
continuous_diffusion_defaults().keys())))
assert args.mixed_prediction, 'enable mixed_prediction by default'
logger.log('create VPSDE diffusion.')
else:
sde_diffusion = None
# if 'cldm' in args.trainer_name:
# assert isinstance(denoise_model, tuple)
# denoise_model, controlNet = denoise_model
# controlNet.to(dist_util.dev())
# controlNet.train()
# else:
# controlNet = None
training_loop_class = TrainLoop(rec_model=auto_encoder,
denoise_model=denoise_model,
control_model=controlNet,
diffusion=diffusion,
sde_diffusion=sde_diffusion,
loss_class=None,
data=None,
eval_data=eval_data,
**vars(args))
logger.log("sampling...")
dist_util.synchronize()
# all_images = []
# all_labels = []
# while len(all_images) * args.batch_size < args.num_samples:
if dist_util.get_rank() == 0:
(Path(logger.get_dir()) / 'FID_Cals').mkdir(exist_ok=True,
parents=True)
with open(os.path.join(args.logdir, 'args.json'), 'w') as f:
json.dump(vars(args), f, indent=2)
# load eval pose
if args.cfg == 'ffhq':
camera = th.load('assets/ffhq_eval_pose.pt',
map_location=dist_util.dev())[:]
elif args.cfg == 'shapenet':
camera = th.load('assets/shapenet_eval_pose.pt',
map_location=dist_util.dev())[:]
for sample_idx in trange(args.num_samples):
model_kwargs = {}
# if args.class_cond:
# classes = th.randint(low=0,
# high=NUM_CLASSES,
# size=(args.batch_size, ),
# device=dist_util.dev())
# model_kwargs["y"] = classes
training_loop_class.step = sample_idx # save to different position
if args.create_controlnet or 'crossattn' in args.trainer_name:
training_loop_class.eval_cldm(
prompt=args.prompt,
unconditional_guidance_scale=args.
unconditional_guidance_scale,
use_ddim=args.use_ddim,
save_img=args.save_img,
use_train_trajectory=args.use_train_trajectory,
export_mesh=args.export_mesh,
camera=camera,
overwrite_diff_inp_size=args.overwrite_diff_inp_size,
# training_loop_class.rec_model,
# training_loop_class.ddpm_model
)
else:
# evaluate ldm
training_loop_class.eval_ddpm_sample(
training_loop_class.rec_model,
save_img=args.save_img,
use_train_trajectory=args.use_train_trajectory,
export_mesh=args.export_mesh,
camera=camera,
# training_loop_class.ddpm_model
)
dist.barrier()
logger.log("sampling complete")
def create_argparser():
defaults = dict(
image_size_encoder=224,
triplane_scaling_divider=1.0, # divide by this value
diffusion_input_size=-1,
trainer_name='adm',
use_amp=False,
# triplane_scaling_divider=1.0, # divide by this value
# * sampling flags
clip_denoised=False,
num_samples=10,
use_ddim=False,
ddpm_model_path="",
cldm_model_path="",
rec_model_path="",
# * eval logging flags
logdir="/mnt/lustre/yslan/logs/nips23/",
data_dir="",
eval_data_dir="",
eval_batch_size=1,
num_workers=1,
# * training flags for loading TrainingLoop class
overfitting=False,
image_size=128,
iterations=150000,
schedule_sampler="uniform",
anneal_lr=False,
lr=5e-5,
weight_decay=0.0,
lr_anneal_steps=0,
batch_size=1,
microbatch=-1, # -1 disables microbatches
ema_rate="0.9999", # comma-separated list of EMA values
log_interval=50,
eval_interval=2500,
save_interval=10000,
resume_checkpoint="",
resume_cldm_checkpoint="",
resume_checkpoint_EG3D="",
use_fp16=False,
fp16_scale_growth=1e-3,
load_submodule_name='', # for loading pretrained auto_encoder model
ignore_resume_opt=False,
freeze_ae=False,
denoised_ae=True,
# inference prompt
prompt="a red chair",
interval=1,
objv_dataset=False,
use_lmdb=False,
save_img=False,
use_train_trajectory=
False, # use train trajectory to sample images for fid calculation
unconditional_guidance_scale=1.0,
cond_key='img_sr',
use_eos_feature=False,
export_mesh=False,
overwrite_diff_inp_size=None,
allow_tf32=True,
)
defaults.update(model_and_diffusion_defaults())
defaults.update(encoder_and_nsr_defaults()) # type: ignore
defaults.update(loss_defaults())
defaults.update(continuous_diffusion_defaults())
defaults.update(control_net_defaults())
defaults.update(dataset_defaults())
parser = argparse.ArgumentParser()
add_dict_to_argparser(parser, defaults)
return parser
if __name__ == "__main__":
# os.environ["TORCH_CPP_LOG_LEVEL"] = "INFO"
# os.environ["NCCL_DEBUG"] = "INFO"
os.environ[
"TORCH_DISTRIBUTED_DEBUG"] = "DETAIL" # set to DETAIL for runtime logging.
args = create_argparser().parse_args()
args.local_rank = int(os.environ["LOCAL_RANK"])
args.gpus = th.cuda.device_count()
args.rendering_kwargs = rendering_options_defaults(args)
main(args)