add gradio example

gradio_app.py

@@ -0,0 +1,237 @@
+import torch
+import argparse
+
+from nerf.provider import NeRFDataset
+from nerf.utils import *
+
+import gradio as gr
+import gc
+
+print(f'[INFO] loading options..')
+
+# fake config object, this should not be used in CMD, only allow change from gradio UI.
+parser = argparse.ArgumentParser()
+parser.add_argument('--text', default=None, help="text prompt")
+# parser.add_argument('-O', action='store_true', help="equals --fp16 --cuda_ray --dir_text")
+# parser.add_argument('-O2', action='store_true', help="equals --fp16 --dir_text")
+parser.add_argument('--test', action='store_true', help="test mode")
+parser.add_argument('--save_mesh', action='store_true', help="export an obj mesh with texture")
+parser.add_argument('--eval_interval', type=int, default=10, help="evaluate on the valid set every interval epochs")
+parser.add_argument('--workspace', type=str, default='trial_gradio')
+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=10000, 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=64, help="num steps sampled per ray (only valid when not using --cuda_ray)")
+parser.add_argument('--upsample_steps', type=int, default=64, 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=1000, help="training iters that only use albedo shading")
+# 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, decrease this if CUDA OOM.
+parser.add_argument('--w', type=int, default=64, help="render width for NeRF in training")
+parser.add_argument('--h', type=int, default=64, help="render height for NeRF in training")
+parser.add_argument('--jitter_pose', action='store_true', help="add jitters to the randomly sampled camera poses")
+
+### 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-encode the text prompt, by appending front/side/back/overhead view")
+parser.add_argument('--angle_overhead', type=float, default=30, help="[0, angle_overhead] is the overhead region")
+parser.add_argument('--angle_front', type=float, default=60, help="[0, angle_front] is the front region, [180, 180+angle_front] the back region, otherwise the side region.")
+
+parser.add_argument('--lambda_entropy', type=float, default=1e-4, help="loss scale for alpha entropy")
+parser.add_argument('--lambda_opacity', type=float, default=0, help="loss scale for alpha value")
+parser.add_argument('--lambda_orient', type=float, default=1e-2, help="loss scale for orientation")
+
+### 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 in [0, 180], corresponding to elevation [90, -90]")
+parser.add_argument('--light_phi', type=float, default=0, help="default GUI light direction in [0, 360), azimuth")
+parser.add_argument('--max_spp', type=int, default=1, help="GUI rendering max sample per pixel")
+
+opt = parser.parse_args() 
+
+# default to use -O !!!
+opt.fp16 = True
+opt.dir_text = True
+opt.cuda_ray = True
+# opt.lambda_entropy = 1e-4
+# opt.lambda_opacity = 0
+
+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)
+
+device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+
+print(f'[INFO] loading models..')
+
+if opt.guidance == 'stable-diffusion':
+    from nerf.sd import StableDiffusion
+    guidance = StableDiffusion(device)
+elif opt.guidance == 'clip':
+    from nerf.clip import CLIP
+    guidance = CLIP(device)
+else:
+    raise NotImplementedError(f'--guidance {opt.guidance} is not implemented.')
+
+train_loader = NeRFDataset(opt, device=device, type='train', H=opt.h, W=opt.w, size=100).dataloader()
+valid_loader = NeRFDataset(opt, device=device, type='val', H=opt.H, W=opt.W, size=5).dataloader()
+test_loader = NeRFDataset(opt, device=device, type='test', H=opt.H, W=opt.W, size=100).dataloader()
+
+print(f'[INFO] everything loaded!')
+
+trainer = None
+model = None
+
+def reset_params(m):
+
+    @torch.no_grad()
+    def _reset(m: nn.Module):
+        reset_parameters = getattr(m, "reset_parameters", None)
+        if callable(reset_parameters):
+            m.reset_parameters()
+
+    model.apply(fn=_reset)
+
+# define UI
+
+with gr.Blocks(css=".gradio-container {max-width: 512px;}") as demo:
+
+    # title
+    gr.Markdown('[Stable-DreamFusion](https://github.com/ashawkey/stable-dreamfusion) Text-to-3D Example')
+
+    # inputs
+    prompt = gr.Textbox(label="Prompt", max_lines=1, value="a DSLR photo of a koi fish")
+    iters = gr.Slider(label="Iters", minimum=1000, maximum=20000, value=5000, step=100)
+    seed = gr.Slider(label="Seed", minimum=0, maximum=2147483647, step=1, randomize=True)
+    button = gr.Button('Generate')
+
+    # outputs
+    image = gr.Image(label="image", visible=True)
+    video = gr.Video(label="video", visible=False)
+    logs = gr.Textbox(label="logging")
+
+    # gradio main func
+    def submit(text, iters, seed):
+
+        global trainer, model
+
+        # seed
+        opt.seed = seed
+        opt.text = text
+        opt.iters = iters
+
+        seed_everything(seed)
+
+        # clean up
+        if trainer is not None:
+            del model
+            del trainer
+            gc.collect()
+            torch.cuda.empty_cache()
+            print('[INFO] clean up!')
+
+        # simply reload everything...
+        model = NeRFNetwork(opt)
+        optimizer = lambda model: torch.optim.Adam(model.get_params(opt.lr), betas=(0.9, 0.99), eps=1e-15)
+        scheduler = lambda optimizer: optim.lr_scheduler.LambdaLR(optimizer, lambda iter: 0.1 ** min(iter / opt.iters, 1))
+
+        trainer = Trainer('df', 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=opt.eval_interval, scheduler_update_every_step=True)
+
+        # train (every ep only contain 8 steps, so we can get some vis every ~10s)
+        STEPS = 8
+        max_epochs = np.ceil(opt.iters / STEPS).astype(np.int32)
+
+        # we have to get the explicit training loop out here to yield progressive results...
+        loader = iter(valid_loader)
+
+        start_t = time.time()
+
+        for epoch in range(max_epochs):
+
+            trainer.train_gui(train_loader, step=STEPS)
+            
+            # manual test and get intermediate results
+            try:
+                data = next(loader)
+            except StopIteration:
+                loader = iter(valid_loader)
+                data = next(loader)
+
+            trainer.model.eval()
+
+            if trainer.ema is not None:
+                trainer.ema.store()
+                trainer.ema.copy_to()
+
+            with torch.no_grad():
+                with torch.cuda.amp.autocast(enabled=trainer.fp16):
+                    preds, preds_depth = trainer.test_step(data, perturb=False)
+
+            if trainer.ema is not None:
+                trainer.ema.restore()
+
+            pred = preds[0].detach().cpu().numpy()
+            # pred_depth = preds_depth[0].detach().cpu().numpy()
+
+            pred = (pred * 255).astype(np.uint8)
+
+            yield {
+                image: gr.update(value=pred, visible=True),
+                video: gr.update(visible=False),
+                logs: f"training iters: {epoch * STEPS} / {iters}, lr: {trainer.optimizer.param_groups[0]['lr']:.6f}",
+            }
+        
+
+        # test
+        trainer.test(test_loader)
+
+        results = glob.glob(os.path.join(opt.workspace, 'results', '*rgb*.mp4'))
+        assert results is not None, "cannot retrieve results!"
+        results.sort(key=lambda x: os.path.getmtime(x)) # sort by mtime
+        
+        end_t = time.time()
+        
+        yield {
+            image: gr.update(visible=False),
+            video: gr.update(value=results[-1], visible=True),
+            logs: f"Generation Finished in {(end_t - start_t)/ 60:.4f} minutes!",
+        }
+
+    
+    button.click(
+        submit, 
+        [prompt, iters, seed],
+        [image, video, logs]
+    )
+
+# concurrency_count: only allow ONE running progress, else GPU will OOM.
+demo.queue(concurrency_count=1)
+
+demo.launch()

main.py

@@ -138,7 +138,7 @@ if __name__ == '__main__':
         scheduler = lambda optimizer: optim.lr_scheduler.LambdaLR(optimizer, lambda iter: 0.1 ** min(iter / opt.iters, 1))
         # scheduler = lambda optimizer: optim.lr_scheduler.OneCycleLR(optimizer, max_lr=opt.lr, total_steps=opt.iters, pct_start=0.1)
 
-        trainer = Trainer('df', 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=opt.eval_interval, scheduler_update_every_step=True)
+        trainer = Trainer('df', opt, model, guidance, device=device, workspace=opt.workspace, optimizer=optimizer, ema_decay=None, fp16=opt.fp16, lr_scheduler=scheduler, use_checkpoint=opt.ckpt, eval_interval=opt.eval_interval, scheduler_update_every_step=True)
 
         if opt.gui:
             trainer.train_loader = train_loader # attach dataloader to trainer

nerf/utils.py

@@ -195,9 +195,6 @@ class Trainer(object):
         self.scheduler_update_every_step = scheduler_update_every_step
         self.device = device if device is not None else torch.device(f'cuda:{local_rank}' if torch.cuda.is_available() else 'cpu')
         self.console = Console()
-
-        # text prompt
-        ref_text = self.opt.text
     
         model.to(self.device)
         if self.world_size > 1:
@@ -208,20 +205,13 @@ class Trainer(object):
         # guide model
         self.guidance = guidance
 
+        # text prompt
         if self.guidance is not None:
-            assert ref_text is not None, 'Training must provide a text prompt!'
-
+            
             for p in self.guidance.parameters():
                 p.requires_grad = False
 
-            if not self.opt.dir_text:
-                self.text_z = self.guidance.get_text_embeds([ref_text])
-            else:
-                self.text_z = []
-                for d in ['front', 'side', 'back', 'side', 'overhead', 'bottom']:
-                    text = f"{ref_text}, {d} view"
-                    text_z = self.guidance.get_text_embeds([text])
-                    self.text_z.append(text_z)
+            self.prepare_text_embeddings()
         
         else:
             self.text_z = None
@@ -257,7 +247,7 @@ class Trainer(object):
             "results": [], # metrics[0], or valid_loss
             "checkpoints": [], # record path of saved ckpt, to automatically remove old ckpt
             "best_result": None,
-            }
+        }
 
         # auto fix
         if len(metrics) == 0 or self.use_loss_as_metric:
@@ -297,6 +287,23 @@ class Trainer(object):
                 self.log(f"[INFO] Loading {self.use_checkpoint} ...")
                 self.load_checkpoint(self.use_checkpoint)
 
+    # calculate the text embs.
+    def prepare_text_embeddings(self):
+
+        if self.opt.text is None:
+            self.log(f"[WARN] text prompt is not provided.")
+            self.text_z = None
+            return
+
+        if not self.opt.dir_text:
+            self.text_z = self.guidance.get_text_embeds([self.opt.text])
+        else:
+            self.text_z = []
+            for d in ['front', 'side', 'back', 'side', 'overhead', 'bottom']:
+                text = f"{self.opt.text}, {d} view"
+                text_z = self.guidance.get_text_embeds([text])
+                self.text_z.append(text_z)
+
     def __del__(self):
         if self.log_ptr: 
             self.log_ptr.close()
@@ -447,6 +454,9 @@ class Trainer(object):
     ### ------------------------------
 
     def train(self, train_loader, valid_loader, max_epochs):
+
+        assert self.text_z is not None, 'Training must provide a text prompt!'
+
         if self.use_tensorboardX and self.local_rank == 0:
             self.writer = tensorboardX.SummaryWriter(os.path.join(self.workspace, "run", self.name))
 

readme.md

@@ -86,6 +86,12 @@ python main.py --text "a hamburger" --workspace trial -O --albedo_iters 10000 #
 # 2. use a smaller density regularization weight:
 python main.py --text "a hamburger" --workspace trial -O --lambda_entropy 1e-5
 
+# you can also train in a GUI to visualize the training progress:
+python main.py --text "a hamburger" --workspace trial -O --gui
+
+# A Gradio GUI is also possible (with less options):
+python gradio_app.py # open in web browser
+
 ## after the training is finished:
 # test (exporting 360 video)
 python main.py --workspace trial -O --test