update

Author: yccyenchicheng

README.md

@@ -15,7 +15,7 @@ https://user-images.githubusercontent.com/27779063/206553305-e01009f7-3131-4a6b-
 # Installation
 We recommend using [`conda`](https://www.anaconda.com/products/distribution) to install the required python packages. You might need to change the `cudatoolkit` version to match with your GPU driver.
 ```
-conda create -n sdf python=3.8 -y && conda activate sdf 
+conda create -n sdfusion python=3.8 -y && conda activate sdfusion
 conda install pytorch==1.9.0 torchvision==0.10.0 torchaudio==0.9.0 cudatoolkit=11.3 -c pytorch -c conda-forge -y
 conda install -c fvcore -c iopath -c conda-forge fvcore iopath -y
 conda install pytorch3d -c pytorch3d

datasets/snet_mm2shape_dataset.py

@@ -62,7 +62,6 @@ def initialize(self, opt, phase='train', cat='all', res=64):
             render_img_list = [os.path.join(render_img_dir, f) for f in os.listdir(render_img_dir) if '.png' in f]
 
             if not os.path.exists(sdf_path):
-                # import pdb; pdb.set_trace()
                 continue
 
             self.model_list.append(sdf_path)

models/sdfusion_mm_model.py

@@ -411,7 +411,7 @@ def forward(self):
     # check: ddpm.py, log_images(). line 1317~1327
     @torch.no_grad()
     # def inference(self, data, sample=True, ddim_steps=None, ddim_eta=0., quantize_denoised=True, infer_all=False):
-    def naive_inference(self, data, ddim_steps=None, ddim_eta=0., uc_scale=None,
+    def inference(self, data, ddim_steps=None, ddim_eta=0., uc_scale=None,
                   infer_all=False, max_sample=16):
 
         self.switch_eval()
@@ -444,10 +444,6 @@ def naive_inference(self, data, ddim_steps=None, ddim_eta=0., uc_scale=None,
         shape = self.z_shape
 
         # get noise, denoise, and decode with vqvae
-        uc = self.cond_model(self.uc_img).float() # img shape
-        c_img = self.cond_model(self.img).float()
-        B = c_img.shape[0]
-        shape = self.z_shape
         samples, intermediates = self.ddim_sampler.sample(S=ddim_steps,
                                                         batch_size=B,
                                                         shape=shape,
@@ -465,7 +461,7 @@ def naive_inference(self, data, ddim_steps=None, ddim_eta=0., uc_scale=None,
     #         txt_scale=1.0, img_scale=1.0, mask_mode='1', mask_x=False,
     #         mm_cls_free=False,
     #     ):
-    def inference(self, data, mask_mode=None, ddim_steps=None, ddim_eta=0., uc_scale=None, 
+    def mm_inference(self, data, mask_mode=None, ddim_steps=None, ddim_eta=0., uc_scale=None, 
                   txt_scale=1.0, img_scale=1.0, mm_cls_free=False, infer_all=False, max_sample=16):
 
         self.switch_eval()
@@ -639,7 +635,7 @@ def get_current_visuals(self):
         b, c, h, w = self.img_gt.shape
         img_shape = (3, h, w)
         # write text as img
-        self.img_text = self.get_img_text(self.txt, bs=b, img_shape=img_shape)
+        self.img_text = self.write_text_on_img(self.txt, bs=b, img_shape=img_shape)
         self.img_text = rearrange(torch.from_numpy(self.img_text), 'b h w c -> b c h w')
 
         vis_tensor_names = [