import torch import argparse from nerf.provider import NeRFDataset from nerf.utils import * from optimizer import Shampoo from nerf.sd import StableDiffusion from nerf.clip import CLIP from nerf.gui import NeRFGUI # torch.autograd.set_detect_anomaly(True) if __name__ == '__main__': parser = argparse.ArgumentParser() parser.add_argument('--text', help="text prompt") parser.add_argument('-O', action='store_true', help="equals --fp16 --cuda_ray --preload") parser.add_argument('--test', action='store_true', help="test mode") parser.add_argument('--workspace', type=str, default='workspace') parser.add_argument('--guidance', type=str, default='stable-diffusion', help='choose from [stable-diffusion, clip]') parser.add_argument('--seed', type=int, default=0) ### training options parser.add_argument('--iters', type=int, default=15000, help="training iters") parser.add_argument('--lr', type=float, default=1e-3, help="initial learning rate") parser.add_argument('--ckpt', type=str, default='latest') parser.add_argument('--cuda_ray', action='store_true', help="use CUDA raymarching instead of pytorch") parser.add_argument('--max_steps', type=int, default=1024, help="max num steps sampled per ray (only valid when using --cuda_ray)") parser.add_argument('--num_steps', type=int, default=256, help="num steps sampled per ray (only valid when not using --cuda_ray)") parser.add_argument('--upsample_steps', type=int, default=0, help="num steps up-sampled per ray (only valid when not using --cuda_ray)") parser.add_argument('--update_extra_interval', type=int, default=16, help="iter interval to update extra status (only valid when using --cuda_ray)") parser.add_argument('--max_ray_batch', type=int, default=4096, help="batch size of rays at inference to avoid OOM (only valid when not using --cuda_ray)") parser.add_argument('--albedo_iters', type=int, default=15000, help="training iters") # model options parser.add_argument('--bg_radius', type=float, default=1.4, help="if positive, use a background model at sphere(bg_radius)") parser.add_argument('--density_thresh', type=float, default=10, help="threshold for density grid to be occupied") # network backbone parser.add_argument('--fp16', action='store_true', help="use amp mixed precision training") parser.add_argument('--backbone', type=str, default='grid', help="nerf backbone, choose from [grid, tcnn, vanilla]") # rendering resolution in training parser.add_argument('--w', type=int, default=64, help="render width for CLIP training (<=224)") parser.add_argument('--h', type=int, default=64, help="render height for CLIP training (<=224)") ### dataset options parser.add_argument('--bound', type=float, default=1, help="assume the scene is bounded in box(-bound, bound)") parser.add_argument('--dt_gamma', type=float, default=0, help="dt_gamma (>=0) for adaptive ray marching. set to 0 to disable, >0 to accelerate rendering (but usually with worse quality)") parser.add_argument('--min_near', type=float, default=0.1, help="minimum near distance for camera") parser.add_argument('--radius_range', type=float, nargs='*', default=[1.0, 1.5], help="training camera radius range") parser.add_argument('--fovy_range', type=float, nargs='*', default=[40, 70], help="training camera fovy range") parser.add_argument('--dir_text', action='store_true', help="direction encoded text prompt") ### GUI options parser.add_argument('--gui', action='store_true', help="start a GUI") parser.add_argument('--W', type=int, default=800, help="GUI width") parser.add_argument('--H', type=int, default=800, help="GUI height") parser.add_argument('--radius', type=float, default=3, help="default GUI camera radius from center") parser.add_argument('--fovy', type=float, default=60, help="default GUI camera fovy") parser.add_argument('--light_theta', type=float, default=60, help="default GUI light direction") parser.add_argument('--light_phi', type=float, default=0, help="default GUI light direction") parser.add_argument('--max_spp', type=int, default=64, help="GUI rendering max sample per pixel") opt = parser.parse_args() if opt.O: opt.fp16 = True opt.cuda_ray = True opt.dir_text = True if opt.backbone == 'vanilla': from nerf.network import NeRFNetwork elif opt.backbone == 'tcnn': from nerf.network_tcnn import NeRFNetwork elif opt.backbone == 'grid': from nerf.network_grid import NeRFNetwork else: raise NotImplementedError(f'--backbone {opt.backbone} is not implemented!') print(opt) seed_everything(opt.seed) model = NeRFNetwork(opt) print(model) device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') if opt.test: guidance = None # do not load guidance at test trainer = Trainer('ngp', opt, model, guidance, device=device, workspace=opt.workspace, fp16=opt.fp16, use_checkpoint=opt.ckpt) if opt.gui: gui = NeRFGUI(opt, trainer) gui.render() else: test_loader = NeRFDataset(opt, device=device, type='test', H=opt.H, W=opt.W, size=100).dataloader() trainer.test(test_loader) trainer.save_mesh(resolution=256) else: if opt.guidance == 'stable-diffusion': guidance = StableDiffusion(device) elif opt.guidance == 'clip': guidance = CLIP(device) else: raise NotImplementedError(f'--guidance {opt.guidance} is not implemented.') optimizer = lambda model: torch.optim.Adam(model.get_params(opt.lr), betas=(0.9, 0.99), eps=1e-15) # optimizer = lambda model: Shampoo(model.get_params(opt.lr)) train_loader = NeRFDataset(opt, device=device, type='train', H=opt.h, W=opt.w, size=100).dataloader() # decay to 0.01 * init_lr at last iter step scheduler = lambda optimizer: optim.lr_scheduler.LambdaLR(optimizer, lambda iter: 0.01 ** min(iter / opt.iters, 1)) trainer = Trainer('ngp', opt, model, guidance, device=device, workspace=opt.workspace, optimizer=optimizer, ema_decay=0.95, fp16=opt.fp16, lr_scheduler=scheduler, use_checkpoint=opt.ckpt, eval_interval=1) if opt.gui: trainer.train_loader = train_loader # attach dataloader to trainer gui = NeRFGUI(opt, trainer) gui.render() else: valid_loader = NeRFDataset(opt, device=device, type='val', H=opt.H, W=opt.W, size=5).dataloader() max_epoch = np.ceil(opt.iters / len(train_loader)).astype(np.int32) trainer.train(train_loader, valid_loader, max_epoch) # also test test_loader = NeRFDataset(opt, device=device, type='test', H=opt.H, W=opt.W, size=100).dataloader() trainer.test(test_loader) trainer.save_mesh(resolution=256)