add ddim support

app.py

@@ -127,7 +127,7 @@ def plot_points(colors, coords):
     return fig
 
 
-def inference(runner: DemoRunner, cfg: ProjectConfig, rgb, mask_hum, mask_obj, std_coverage, input_seed, input_cls):
+def inference(runner: DemoRunner, cfg: ProjectConfig, rgb, mask_hum, mask_obj, std_coverage, input_seed, input_cls, input_scheduler):
     """
     given user input, run inference
     :param runner:
@@ -138,6 +138,7 @@ def inference(runner: DemoRunner, cfg: ProjectConfig, rgb, mask_hum, mask_obj, s
     :param std_coverage: float value, used to estimate camera translation
     :param input_seed: random seed
     :param input_cls: the object category of the input image
+    :param input_scheduler: reverse sampling scheduler, ddim or ddpm
     :return: path to the 3D reconstruction, and an interactive 3D figure for visualizing the point cloud
     """
     log = ""
@@ -153,6 +154,8 @@ def inference(runner: DemoRunner, cfg: ProjectConfig, rgb, mask_hum, mask_obj, s
             log += f"Reloading fine-tuned checkpoint of category {input_cls}\n"
             runner.reload_checkpoint(input_cls)
 
+        cfg.run.diffusion_scheduler = input_scheduler
+        cfg.run.num_inference_steps = 1000 if input_scheduler == 'ddpm' else 100
         out_stage1, out_stage2 = runner.forward_batch(batch, cfg)
         points = out_stage2.points_packed().cpu().numpy()
         colors = out_stage2.features_packed().cpu().numpy()
@@ -204,6 +207,10 @@ def main(cfg: ProjectConfig):
                                                  'chair', 'skateboard', 'suitcase', 'table'],
                                         value='general')
                 input_seed = gr.Number(label='Random seed', value=42)
+                input_scheduler = gr.Dropdown(label='Diffusion scheduler',
+                                        info='Reverse diffusion scheduler: DDIM is 10x faster',
+                                        choices=['ddpm', 'ddim'],
+                                        value='ddim')
         # Output visualization
         with gr.Row():
             pc_plot = gr.Plot(label="Reconstructed point cloud")
@@ -217,20 +224,20 @@ def main(cfg: ProjectConfig):
         with gr.Row():
             button_recon = gr.Button("Start Reconstruction", interactive=True, variant='secondary')
             button_recon.click(fn=partial(inference, runner, cfg),
-                               inputs=[input_rgb, input_mask_hum, input_mask_obj, input_std, input_seed, input_cls],
+                               inputs=[input_rgb, input_mask_hum, input_mask_obj, input_std, input_seed, input_cls, input_scheduler],
                                outputs=[pc_plot, out_pc_download, out_log])
         gr.HTML("""<br/>""")
         # Example input
         example_dir = cfg.run.code_dir_abs+"/examples"
         rgb, ps, obj = 'k1_color.jpg', 'k1_person_mask.png', 'k1_obj_rend_mask.png'
         example_images = gr.Examples([
-            [f"{example_dir}/017450/{rgb}", f"{example_dir}/017450/{ps}", f"{example_dir}/017450/{obj}", 3.0, 42, 'skateboard'],
-            [f"{example_dir}/205904/{rgb}", f"{example_dir}/205904/{ps}", f"{example_dir}/205904/{obj}", 3.2, 42, 'suitcase'],
-            [f"{example_dir}/066241/{rgb}", f"{example_dir}/066241/{ps}", f"{example_dir}/066241/{obj}", 3.5, 42, 'backpack'],
-            [f"{example_dir}/053431/{rgb}", f"{example_dir}/053431/{ps}", f"{example_dir}/053431/{obj}", 3.8, 42, 'chair'],
-            [f"{example_dir}/158107/{rgb}", f"{example_dir}/158107/{ps}", f"{example_dir}/158107/{obj}", 3.8, 42, 'chair'],
+            [f"{example_dir}/017450/{rgb}", f"{example_dir}/017450/{ps}", f"{example_dir}/017450/{obj}", 3.0, 42, 'skateboard', 'ddim'],
+            [f"{example_dir}/205904/{rgb}", f"{example_dir}/205904/{ps}", f"{example_dir}/205904/{obj}", 3.2, 42, 'suitcase', 'ddim'],
+            [f"{example_dir}/066241/{rgb}", f"{example_dir}/066241/{ps}", f"{example_dir}/066241/{obj}", 3.5, 42, 'backpack', 'ddim'],
+            [f"{example_dir}/053431/{rgb}", f"{example_dir}/053431/{ps}", f"{example_dir}/053431/{obj}", 3.8, 42, 'chair', 'ddim'],
+            [f"{example_dir}/158107/{rgb}", f"{example_dir}/158107/{ps}", f"{example_dir}/158107/{obj}", 3.8, 42, 'chair', 'ddim'],
 
-        ], inputs=[input_rgb, input_mask_hum, input_mask_obj, input_std, input_seed, input_cls],)
+        ], inputs=[input_rgb, input_mask_hum, input_mask_obj, input_std, input_seed, input_cls, input_scheduler],)
 
         gr.Markdown(citation_str)
 

configs/structured.py

@@ -127,6 +127,7 @@ class PointCloudDiffusionModelConfig(PointCloudProjectionModelConfig):
     beta_end: float = 8e-3  # 0.012
     beta_schedule: str = 'linear'  # 'custom'
     dm_pred_type: str = 'epsilon'  # diffusion model prediction type, sample (x0) or noise
+    ddim_eta: float = 1.0  # DDIM eta parameter: 0 is the default one which does deterministic generation
 
     # Point cloud model arguments
     point_cloud_model: str = 'pvcnn'

demo.py

@@ -180,6 +180,7 @@ class DemoRunner:
                                                       mask=torch.stack(batch['masks']).to('cuda'),
                                                       scheduler=cfg.run.diffusion_scheduler,
                                                       num_inference_steps=cfg.run.num_inference_steps,
+                                                      eta=cfg.model.ddim_eta,
                                                       )
         # segment and normalize human/object
         bs = len(out_stage1)
@@ -254,7 +255,11 @@ class DemoRunner:
             radius_hum=radius_hum.unsqueeze(-1),
             radius_obj=radius_obj.unsqueeze(-1),
             sample_from_interm=True,
-            noise_step=cfg.run.sample_noise_step)
+            noise_step=cfg.run.sample_noise_step,
+            scheduler=cfg.run.diffusion_scheduler,
+            num_inference_steps=cfg.run.num_inference_steps,
+            eta=cfg.model.ddim_eta,
+        )
         return out_stage1, out_stage2
 
     def upsample_predicted_pc(self, num_samples, pc_obj):

model/model_hoattn.py

@@ -11,6 +11,7 @@ import numpy as np
 
 from pytorch3d.structures import Pointclouds
 from pytorch3d.renderer import CamerasBase
+from diffusers.schedulers import DDPMScheduler, DDIMScheduler
 from .model_diff_data import ConditionalPCDiffusionBehave
 from .pvcnn.pvcnn_ho import PVCNN2HumObj
 import torch.nn.functional as F
@@ -375,17 +376,30 @@ class CrossAttenHODiffusionModel(ConditionalPCDiffusionBehave):
 
         return (output, all_outputs) if return_all_outputs else output
 
-    def get_reverse_timesteps(self, scheduler, interm_steps:int):
+    def get_reverse_timesteps(self, scheduler, interm_steps: int):
         """
-
+        get the timesteps to run reverse diffusion
         :param scheduler:
-        :param interm_steps: start from some intermediate steps
+        :param interm_steps: start from some intermediate steps, the step number is for DDPM scheduler
+            if DDIM, will be recomputed accordingly
         :return:
         """
-        if interm_steps > 0:
-            timesteps = torch.from_numpy(np.arange(0, interm_steps)[::-1].copy()).to(self.device)
+        if isinstance(scheduler, DDPMScheduler):
+            # DDPM, directly reverse N steps from interm_steps
+            if interm_steps > 0:
+                timesteps = torch.from_numpy(np.arange(0, interm_steps)[::-1].copy()).to(self.device)
+            else:
+                timesteps = scheduler.timesteps.to(self.device)
+        elif isinstance(scheduler, DDIMScheduler):
+            if interm_steps > 0:
+                # compute a step ratio, and find the intermediate steps for DDIM
+                step_ratio = scheduler.config.num_train_timesteps // scheduler.num_inference_steps
+                timesteps = (np.arange(0, interm_steps, step_ratio)).round()[::-1].copy().astype(np.int64)
+                timesteps = torch.from_numpy(timesteps).to(self.device)
+            else:
+                timesteps = scheduler.timesteps.to(self.device)
         else:
-            timesteps = scheduler.timesteps.to(self.device)
+            raise NotImplementedError
         return timesteps
 
     def pack_norm_params(self, kwargs:dict, scale=True):