reformat

Signed-off-by: Can-Zhao <[email protected]>

Author: Can-Zhao

monai/inferers/inferer.py

@@ -1402,12 +1402,18 @@ def sample(  # type: ignore[override]
         if not scheduler:
             scheduler = self.scheduler
         image = input_noise
+
+        all_next_timesteps = torch.cat((scheduler.timesteps[1:], torch.tensor([0], dtype=scheduler.timesteps.dtype)))
         if verbose and has_tqdm:
-            progress_bar = tqdm(scheduler.timesteps)
+            progress_bar = tqdm(
+                zip(scheduler.timesteps, all_next_timesteps),
+                total=min(len(scheduler.timesteps), len(all_next_timesteps)),
+            )
         else:
-            progress_bar = iter(scheduler.timesteps)
+            progress_bar = iter(zip(scheduler.timesteps, all_next_timesteps))
         intermediates = []
-        for t in progress_bar:
+
+        for t, next_t in progress_bar:
             diffuse = diffusion_model
             if isinstance(diffusion_model, SPADEDiffusionModelUNet):
                 diffuse = partial(diffusion_model, seg=seg)
@@ -1446,7 +1452,11 @@ def sample(  # type: ignore[override]
                 )
 
             # 3. compute previous image: x_t -> x_t-1
-            image, _ = scheduler.step(model_output, t, image)  # type: ignore[operator]
+            if not isinstance(scheduler, RFlowScheduler):
+                image, _ = scheduler.step(model_output, t, image)
+            else:
+                image, _ = scheduler.step(model_output, t, image, next_t)
+
             if save_intermediates and t % intermediate_steps == 0:
                 intermediates.append(image)
         if save_intermediates:

tests/inferers/test_controlnet_inferers.py

@@ -587,7 +587,7 @@ def test_sampler_conditioned(self, model_params, controlnet_params, input_shape)
         controlnet.eval()
         mask = torch.randn(input_shape).to(device)
         noise = torch.randn(input_shape).to(device)
-        
+
         # DDIM
         scheduler = DDIMScheduler(num_train_timesteps=1000)
         inferer = ControlNetDiffusionInferer(scheduler=scheduler)
@@ -755,11 +755,11 @@ def test_prediction_shape(
         input = torch.randn(input_shape).to(device)
         mask = torch.randn(input_shape).to(device)
         noise = torch.randn(latent_shape).to(device)
+
         for scheduler in [DDPMScheduler(num_train_timesteps=10), RFlowScheduler(num_train_timesteps=1000)]:
             inferer = ControlNetLatentDiffusionInferer(scheduler=scheduler, scale_factor=1.0)
             scheduler.set_timesteps(num_inference_steps=10)
             timesteps = torch.randint(0, scheduler.num_train_timesteps, (input_shape[0],), device=input.device).long()
-    
             if dm_model_type == "SPADEDiffusionModelUNet":
                 input_shape_seg = list(input_shape)
                 if "label_nc" in stage_2_params.keys():

tests/networks/schedulers/test_scheduler_rflow.py

@@ -21,18 +21,34 @@
 
 TEST_2D_CASE = []
 for sample_method in ["uniform", "logit-normal"]:
-    TEST_2D_CASE.append([{"sample_method": sample_method, "use_timestep_transform": False}, (2, 6, 16, 16), (2, 6, 16, 16)])
+    TEST_2D_CASE.append(
+        [{"sample_method": sample_method, "use_timestep_transform": False}, (2, 6, 16, 16), (2, 6, 16, 16)]
+    )
 
 for sample_method in ["uniform", "logit-normal"]:
-    TEST_2D_CASE.append([{"sample_method": sample_method, "use_timestep_transform": True, "spatial_dim": 2}, (2, 6, 16, 16), (2, 6, 16, 16)])
+    TEST_2D_CASE.append(
+        [
+            {"sample_method": sample_method, "use_timestep_transform": True, "spatial_dim": 2},
+            (2, 6, 16, 16),
+            (2, 6, 16, 16),
+        ]
+    )
 
 
 TEST_3D_CASE = []
 for sample_method in ["uniform", "logit-normal"]:
-    TEST_3D_CASE.append([{"sample_method": sample_method, "use_timestep_transform": False}, (2, 6, 16, 16, 16), (2, 6, 16, 16, 16)])
+    TEST_3D_CASE.append(
+        [{"sample_method": sample_method, "use_timestep_transform": False}, (2, 6, 16, 16, 16), (2, 6, 16, 16, 16)]
+    )
 
 for sample_method in ["uniform", "logit-normal"]:
-    TEST_3D_CASE.append([{"sample_method": sample_method, "use_timestep_transform": True, "spatial_dim": 3}, (2, 6, 16, 16, 16), (2, 6, 16, 16, 16)])
+    TEST_3D_CASE.append(
+        [
+            {"sample_method": sample_method, "use_timestep_transform": True, "spatial_dim": 3},
+            (2, 6, 16, 16, 16),
+            (2, 6, 16, 16, 16),
+        ]
+    )
 
 TEST_CASES = TEST_2D_CASE + TEST_3D_CASE
 
@@ -54,35 +70,35 @@ def test_add_noise(self, input_param, input_shape, expected_shape):
 
     @parameterized.expand(TEST_CASES)
     def test_step_shape(self, input_param, input_shape, expected_shape):
-        scheduler = RFlowScheduler(**input_param)        
+        scheduler = RFlowScheduler(**input_param)
         model_output = torch.randn(input_shape)
         sample = torch.randn(input_shape)
-        scheduler.set_timesteps(num_inference_steps=100, input_img_size_numel=torch.numel(sample[0,0,...]))
+        scheduler.set_timesteps(num_inference_steps=100, input_img_size_numel=torch.numel(sample[0, 0, ...]))
         output_step = scheduler.step(model_output=model_output, timestep=500, sample=sample)
         self.assertEqual(output_step[0].shape, expected_shape)
         self.assertEqual(output_step[1].shape, expected_shape)
 
     @parameterized.expand(TEST_FULl_LOOP)
     def test_full_timestep_loop(self, input_param, input_shape, expected_output):
-        scheduler = RFlowScheduler(**input_param)        
+        scheduler = RFlowScheduler(**input_param)
         torch.manual_seed(42)
         model_output = torch.randn(input_shape)
         sample = torch.randn(input_shape)
-        scheduler.set_timesteps(50, input_img_size_numel=torch.numel(sample[0,0,...]))
+        scheduler.set_timesteps(50, input_img_size_numel=torch.numel(sample[0, 0, ...]))
         for t in range(50):
             sample, _ = scheduler.step(model_output=model_output, timestep=t, sample=sample)
         assert_allclose(sample, expected_output, rtol=1e-3, atol=1e-3)
 
     def test_set_timesteps(self):
         scheduler = RFlowScheduler(num_train_timesteps=1000)
-        scheduler.set_timesteps(num_inference_steps=100, input_img_size_numel=16*16*16)
+        scheduler.set_timesteps(num_inference_steps=100, input_img_size_numel=16 * 16 * 16)
         self.assertEqual(scheduler.num_inference_steps, 100)
         self.assertEqual(len(scheduler.timesteps), 100)
 
     def test_set_timesteps_with_num_inference_steps_bigger_than_num_train_timesteps(self):
         scheduler = RFlowScheduler(num_train_timesteps=1000)
         with self.assertRaises(ValueError):
-            scheduler.set_timesteps(num_inference_steps=2000, input_img_size_numel=16*16*16)
+            scheduler.set_timesteps(num_inference_steps=2000, input_img_size_numel=16 * 16 * 16)
 
 
 if __name__ == "__main__":