Add rectified flow noise scheduler for accelerated diffusion model (#8374)

Fixes # .

### Description

A few sentences describing the changes proposed in this pull request.

### Types of changes
<!--- Put an `x` in all the boxes that apply, and remove the not
applicable items -->
- [x] Non-breaking change (fix or new feature that would not break
existing functionality).
- [ ] Breaking change (fix or new feature that would cause existing
functionality to change).
- [x] New tests added to cover the changes.
- [x] Integration tests passed locally by running `./runtests.sh -f -u
--net --coverage`.
- [ ] Quick tests passed locally by running `./runtests.sh --quick
--unittests --disttests`.
- [x] In-line docstrings updated.
- [ ] Documentation updated, tested `make html` command in the `docs/`
folder.

---------

Signed-off-by: Yiheng Wang <[email protected]>
Signed-off-by: Can-Zhao <[email protected]>
Signed-off-by: pooya-mohammadi <[email protected]>
Signed-off-by: YunLiu <[email protected]>
Signed-off-by: advcu987 <[email protected]>
Signed-off-by: advcu <[email protected]>
Signed-off-by: Eric Kerfoot <[email protected]>
Signed-off-by: Can-Zhao <[email protected]>
Signed-off-by: R. Garcia-Dias <[email protected]>
Signed-off-by: Rafael Garcia-Dias <[email protected]>
Signed-off-by: Eric Kerfoot <[email protected]>
Signed-off-by: James Butler <[email protected]>
Signed-off-by: Virginia Fernandez <[email protected]>
Signed-off-by: Nicolas Kaenzig <[email protected]>
Signed-off-by: Bartosz Grabowski <[email protected]>
Signed-off-by: thibaultdvx <[email protected]>
Signed-off-by: Thibault de Varax <[email protected]>
Signed-off-by: monai-bot <[email protected]>
Co-authored-by: Yiheng Wang <[email protected]>
Co-authored-by: Pooya Mohammadi Kazaj <[email protected]>
Co-authored-by: YunLiu <[email protected]>
Co-authored-by: Eric Kerfoot <[email protected]>
Co-authored-by: advcu <[email protected]>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Rafael Garcia-Dias <[email protected]>
Co-authored-by: Rafael Garcia-Dias <[email protected]>
Co-authored-by: Ben Murray <[email protected]>
Co-authored-by: James Butler <[email protected]>
Co-authored-by: Virginia Fernandez <[email protected]>
Co-authored-by: Virginia Fernandez <[email protected]>
Co-authored-by: Nicolas Känzig <[email protected]>
Co-authored-by: Bartosz Grabowski <[email protected]>
Co-authored-by: Thibault de Varax <[email protected]>
Co-authored-by: monai-bot <[email protected]>

Author: 17 people

docs/source/networks.rst

@@ -750,3 +750,38 @@ Utilities
 
 .. automodule:: monai.apps.reconstruction.networks.nets.utils
   :members:
+
+Noise Schedulers
+----------------
+.. automodule:: monai.networks.schedulers
+.. currentmodule:: monai.networks.schedulers
+
+`Scheduler`
+~~~~~~~~~~~
+.. autoclass:: Scheduler
+  :members:
+
+`NoiseSchedules`
+~~~~~~~~~~~~~~~~
+.. autoclass:: NoiseSchedules
+  :members:
+
+`DDPMScheduler`
+~~~~~~~~~~~~~~~
+.. autoclass:: DDPMScheduler
+  :members:
+
+`DDIMScheduler`
+~~~~~~~~~~~~~~~
+.. autoclass:: DDIMScheduler
+  :members:
+
+`PNDMScheduler`
+~~~~~~~~~~~~~~~
+.. autoclass:: PNDMScheduler
+  :members:
+
+`RFlowScheduler`
+~~~~~~~~~~~~~~~~
+.. autoclass:: RFlowScheduler
+  :members:

monai/apps/generation/maisi/networks/autoencoderkl_maisi.py

@@ -232,6 +232,10 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
         if self.print_info:
             logger.info(f"Number of splits: {self.num_splits}")
 
+        if self.dim_split <= 1 and self.num_splits <= 1:
+            x = self.conv(x)
+            return x
+
         # compute size of splits
         l = x.size(self.dim_split + 2)
         split_size = l // self.num_splits

monai/inferers/inferer.py

@@ -39,7 +39,7 @@
     SPADEAutoencoderKL,
     SPADEDiffusionModelUNet,
 )
-from monai.networks.schedulers import Scheduler
+from monai.networks.schedulers import RFlowScheduler, Scheduler
 from monai.transforms import CenterSpatialCrop, SpatialPad
 from monai.utils import BlendMode, Ordering, PatchKeys, PytorchPadMode, ensure_tuple, optional_import
 from monai.visualize import CAM, GradCAM, GradCAMpp
@@ -859,12 +859,18 @@ def sample(
         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:
             # 1. predict noise model_output
             diffusion_model = (
                 partial(diffusion_model, seg=seg)
@@ -882,9 +888,13 @@ def sample(
                 )
 
             # 2. 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)  # type: ignore
+            else:
+                image, _ = scheduler.step(model_output, t, image, next_t)  # type: ignore
             if save_intermediates and t % intermediate_steps == 0:
                 intermediates.append(image)
+
         if save_intermediates:
             return image, intermediates
         else:
@@ -1392,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)
@@ -1436,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)  # type: ignore
+            else:
+                image, _ = scheduler.step(model_output, t, image, next_t)  # type: ignore
+
             if save_intermediates and t % intermediate_steps == 0:
                 intermediates.append(image)
         if save_intermediates:

monai/networks/schedulers/__init__.py

@@ -14,4 +14,5 @@
 from .ddim import DDIMScheduler
 from .ddpm import DDPMScheduler
 from .pndm import PNDMScheduler
+from .rectified_flow import RFlowScheduler
 from .scheduler import NoiseSchedules, Scheduler