monai-weekly 1.5.dev2510__py3-none-any.whl → 1.5.dev2512__py3-none-any.whl

tests/inferers/test_latent_diffusion_inferer.py

@@ -19,7 +19,7 @@ from parameterized import parameterized
 
 from monai.inferers import LatentDiffusionInferer
 from monai.networks.nets import VQVAE, AutoencoderKL, DiffusionModelUNet, SPADEAutoencoderKL, SPADEDiffusionModelUNet
-from monai.networks.schedulers import DDPMScheduler
+from monai.networks.schedulers import DDPMScheduler, RFlowScheduler
 from monai.utils import optional_import
 
 _, has_einops = optional_import("einops")
@@ -339,31 +339,32 @@ class TestDiffusionSamplingInferer(unittest.TestCase):
 
         input = torch.randn(input_shape).to(device)
         noise = torch.randn(latent_shape).to(device)
-        scheduler = DDPMScheduler(num_train_timesteps=10)
-        inferer = LatentDiffusionInferer(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():
-                input_shape_seg[1] = stage_2_params["label_nc"]
+        for scheduler in [DDPMScheduler(num_train_timesteps=10), RFlowScheduler(num_train_timesteps=1000)]:
+            inferer = LatentDiffusionInferer(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():
+                    input_shape_seg[1] = stage_2_params["label_nc"]
+                else:
+                    input_shape_seg[1] = autoencoder_params["label_nc"]
+                input_seg = torch.randn(input_shape_seg).to(device)
+                prediction = inferer(
+                    inputs=input,
+                    autoencoder_model=stage_1,
+                    diffusion_model=stage_2,
+                    seg=input_seg,
+                    noise=noise,
+                    timesteps=timesteps,
+                )
             else:
-                input_shape_seg[1] = autoencoder_params["label_nc"]
-            input_seg = torch.randn(input_shape_seg).to(device)
-            prediction = inferer(
-                inputs=input,
-                autoencoder_model=stage_1,
-                diffusion_model=stage_2,
-                seg=input_seg,
-                noise=noise,
-                timesteps=timesteps,
-            )
-        else:
-            prediction = inferer(
-                inputs=input, autoencoder_model=stage_1, diffusion_model=stage_2, noise=noise, timesteps=timesteps
-            )
-        self.assertEqual(prediction.shape, latent_shape)
+                prediction = inferer(
+                    inputs=input, autoencoder_model=stage_1, diffusion_model=stage_2, noise=noise, timesteps=timesteps
+                )
+            self.assertEqual(prediction.shape, latent_shape)
 
     @parameterized.expand(TEST_CASES)
     @skipUnless(has_einops, "Requires einops")
@@ -388,29 +389,30 @@ class TestDiffusionSamplingInferer(unittest.TestCase):
         stage_2.eval()
 
         noise = torch.randn(latent_shape).to(device)
-        scheduler = DDPMScheduler(num_train_timesteps=10)
-        inferer = LatentDiffusionInferer(scheduler=scheduler, scale_factor=1.0)
-        scheduler.set_timesteps(num_inference_steps=10)
 
-        if ae_model_type == "SPADEAutoencoderKL" or dm_model_type == "SPADEDiffusionModelUNet":
-            input_shape_seg = list(input_shape)
-            if "label_nc" in stage_2_params.keys():
-                input_shape_seg[1] = stage_2_params["label_nc"]
+        for scheduler in [DDPMScheduler(num_train_timesteps=10), RFlowScheduler(num_train_timesteps=1000)]:
+            inferer = LatentDiffusionInferer(scheduler=scheduler, scale_factor=1.0)
+            scheduler.set_timesteps(num_inference_steps=10)
+
+            if ae_model_type == "SPADEAutoencoderKL" or dm_model_type == "SPADEDiffusionModelUNet":
+                input_shape_seg = list(input_shape)
+                if "label_nc" in stage_2_params.keys():
+                    input_shape_seg[1] = stage_2_params["label_nc"]
+                else:
+                    input_shape_seg[1] = autoencoder_params["label_nc"]
+                input_seg = torch.randn(input_shape_seg).to(device)
+                sample = inferer.sample(
+                    input_noise=noise,
+                    autoencoder_model=stage_1,
+                    diffusion_model=stage_2,
+                    scheduler=scheduler,
+                    seg=input_seg,
+                )
             else:
-                input_shape_seg[1] = autoencoder_params["label_nc"]
-            input_seg = torch.randn(input_shape_seg).to(device)
-            sample = inferer.sample(
-                input_noise=noise,
-                autoencoder_model=stage_1,
-                diffusion_model=stage_2,
-                scheduler=scheduler,
-                seg=input_seg,
-            )
-        else:
-            sample = inferer.sample(
-                input_noise=noise, autoencoder_model=stage_1, diffusion_model=stage_2, scheduler=scheduler
-            )
-        self.assertEqual(sample.shape, input_shape)
+                sample = inferer.sample(
+                    input_noise=noise, autoencoder_model=stage_1, diffusion_model=stage_2, scheduler=scheduler
+                )
+            self.assertEqual(sample.shape, input_shape)
 
     @parameterized.expand(TEST_CASES)
     @skipUnless(has_einops, "Requires einops")
@@ -437,37 +439,38 @@ class TestDiffusionSamplingInferer(unittest.TestCase):
         stage_2.eval()
 
         noise = torch.randn(latent_shape).to(device)
-        scheduler = DDPMScheduler(num_train_timesteps=10)
-        inferer = LatentDiffusionInferer(scheduler=scheduler, scale_factor=1.0)
-        scheduler.set_timesteps(num_inference_steps=10)
 
-        if ae_model_type == "SPADEAutoencoderKL" or dm_model_type == "SPADEDiffusionModelUNet":
-            input_shape_seg = list(input_shape)
-            if "label_nc" in stage_2_params.keys():
-                input_shape_seg[1] = stage_2_params["label_nc"]
+        for scheduler in [DDPMScheduler(num_train_timesteps=10), RFlowScheduler(num_train_timesteps=1000)]:
+            inferer = LatentDiffusionInferer(scheduler=scheduler, scale_factor=1.0)
+            scheduler.set_timesteps(num_inference_steps=10)
+
+            if ae_model_type == "SPADEAutoencoderKL" or dm_model_type == "SPADEDiffusionModelUNet":
+                input_shape_seg = list(input_shape)
+                if "label_nc" in stage_2_params.keys():
+                    input_shape_seg[1] = stage_2_params["label_nc"]
+                else:
+                    input_shape_seg[1] = autoencoder_params["label_nc"]
+                input_seg = torch.randn(input_shape_seg).to(device)
+                sample, intermediates = inferer.sample(
+                    input_noise=noise,
+                    autoencoder_model=stage_1,
+                    diffusion_model=stage_2,
+                    scheduler=scheduler,
+                    seg=input_seg,
+                    save_intermediates=True,
+                    intermediate_steps=1,
+                )
             else:
-                input_shape_seg[1] = autoencoder_params["label_nc"]
-            input_seg = torch.randn(input_shape_seg).to(device)
-            sample, intermediates = inferer.sample(
-                input_noise=noise,
-                autoencoder_model=stage_1,
-                diffusion_model=stage_2,
-                scheduler=scheduler,
-                seg=input_seg,
-                save_intermediates=True,
-                intermediate_steps=1,
-            )
-        else:
-            sample, intermediates = inferer.sample(
-                input_noise=noise,
-                autoencoder_model=stage_1,
-                diffusion_model=stage_2,
-                scheduler=scheduler,
-                save_intermediates=True,
-                intermediate_steps=1,
-            )
-        self.assertEqual(len(intermediates), 10)
-        self.assertEqual(intermediates[0].shape, input_shape)
+                sample, intermediates = inferer.sample(
+                    input_noise=noise,
+                    autoencoder_model=stage_1,
+                    diffusion_model=stage_2,
+                    scheduler=scheduler,
+                    save_intermediates=True,
+                    intermediate_steps=1,
+                )
+            self.assertEqual(len(intermediates), 10)
+            self.assertEqual(intermediates[0].shape, input_shape)
 
     @parameterized.expand(TEST_CASES)
     @skipUnless(has_einops, "Requires einops")
@@ -614,40 +617,40 @@ class TestDiffusionSamplingInferer(unittest.TestCase):
         conditioning_shape[1] = n_concat_channel
         conditioning = torch.randn(conditioning_shape).to(device)
 
-        scheduler = DDPMScheduler(num_train_timesteps=10)
-        inferer = LatentDiffusionInferer(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():
-                input_shape_seg[1] = stage_2_params["label_nc"]
+        for scheduler in [DDPMScheduler(num_train_timesteps=10), RFlowScheduler(num_train_timesteps=1000)]:
+            inferer = LatentDiffusionInferer(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():
+                    input_shape_seg[1] = stage_2_params["label_nc"]
+                else:
+                    input_shape_seg[1] = autoencoder_params["label_nc"]
+                input_seg = torch.randn(input_shape_seg).to(device)
+                prediction = inferer(
+                    inputs=input,
+                    autoencoder_model=stage_1,
+                    diffusion_model=stage_2,
+                    noise=noise,
+                    timesteps=timesteps,
+                    condition=conditioning,
+                    mode="concat",
+                    seg=input_seg,
+                )
             else:
-                input_shape_seg[1] = autoencoder_params["label_nc"]
-            input_seg = torch.randn(input_shape_seg).to(device)
-            prediction = inferer(
-                inputs=input,
-                autoencoder_model=stage_1,
-                diffusion_model=stage_2,
-                noise=noise,
-                timesteps=timesteps,
-                condition=conditioning,
-                mode="concat",
-                seg=input_seg,
-            )
-        else:
-            prediction = inferer(
-                inputs=input,
-                autoencoder_model=stage_1,
-                diffusion_model=stage_2,
-                noise=noise,
-                timesteps=timesteps,
-                condition=conditioning,
-                mode="concat",
-            )
-        self.assertEqual(prediction.shape, latent_shape)
+                prediction = inferer(
+                    inputs=input,
+                    autoencoder_model=stage_1,
+                    diffusion_model=stage_2,
+                    noise=noise,
+                    timesteps=timesteps,
+                    condition=conditioning,
+                    mode="concat",
+                )
+            self.assertEqual(prediction.shape, latent_shape)
 
     @parameterized.expand(TEST_CASES)
     @skipUnless(has_einops, "Requires einops")
@@ -681,36 +684,36 @@ class TestDiffusionSamplingInferer(unittest.TestCase):
         conditioning_shape[1] = n_concat_channel
         conditioning = torch.randn(conditioning_shape).to(device)
 
-        scheduler = DDPMScheduler(num_train_timesteps=10)
-        inferer = LatentDiffusionInferer(scheduler=scheduler, scale_factor=1.0)
-        scheduler.set_timesteps(num_inference_steps=10)
-
-        if dm_model_type == "SPADEDiffusionModelUNet":
-            input_shape_seg = list(input_shape)
-            if "label_nc" in stage_2_params.keys():
-                input_shape_seg[1] = stage_2_params["label_nc"]
+        for scheduler in [DDPMScheduler(num_train_timesteps=10), RFlowScheduler(num_train_timesteps=1000)]:
+            inferer = LatentDiffusionInferer(scheduler=scheduler, scale_factor=1.0)
+            scheduler.set_timesteps(num_inference_steps=10)
+
+            if dm_model_type == "SPADEDiffusionModelUNet":
+                input_shape_seg = list(input_shape)
+                if "label_nc" in stage_2_params.keys():
+                    input_shape_seg[1] = stage_2_params["label_nc"]
+                else:
+                    input_shape_seg[1] = autoencoder_params["label_nc"]
+                input_seg = torch.randn(input_shape_seg).to(device)
+                sample = inferer.sample(
+                    input_noise=noise,
+                    autoencoder_model=stage_1,
+                    diffusion_model=stage_2,
+                    scheduler=scheduler,
+                    conditioning=conditioning,
+                    mode="concat",
+                    seg=input_seg,
+                )
             else:
-                input_shape_seg[1] = autoencoder_params["label_nc"]
-            input_seg = torch.randn(input_shape_seg).to(device)
-            sample = inferer.sample(
-                input_noise=noise,
-                autoencoder_model=stage_1,
-                diffusion_model=stage_2,
-                scheduler=scheduler,
-                conditioning=conditioning,
-                mode="concat",
-                seg=input_seg,
-            )
-        else:
-            sample = inferer.sample(
-                input_noise=noise,
-                autoencoder_model=stage_1,
-                diffusion_model=stage_2,
-                scheduler=scheduler,
-                conditioning=conditioning,
-                mode="concat",
-            )
-        self.assertEqual(sample.shape, input_shape)
+                sample = inferer.sample(
+                    input_noise=noise,
+                    autoencoder_model=stage_1,
+                    diffusion_model=stage_2,
+                    scheduler=scheduler,
+                    conditioning=conditioning,
+                    mode="concat",
+                )
+            self.assertEqual(sample.shape, input_shape)
 
     @parameterized.expand(TEST_CASES_DIFF_SHAPES)
     @skipUnless(has_einops, "Requires einops")
@@ -738,39 +741,39 @@ class TestDiffusionSamplingInferer(unittest.TestCase):
 
         input = torch.randn(input_shape).to(device)
         noise = torch.randn(latent_shape).to(device)
-        scheduler = DDPMScheduler(num_train_timesteps=10)
-        # We infer the VAE shape
-        autoencoder_latent_shape = [i // (2 ** (len(autoencoder_params["channels"]) - 1)) for i in input_shape[2:]]
-        inferer = LatentDiffusionInferer(
-            scheduler=scheduler,
-            scale_factor=1.0,
-            ldm_latent_shape=list(latent_shape[2:]),
-            autoencoder_latent_shape=autoencoder_latent_shape,
-        )
-        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():
-                input_shape_seg[1] = stage_2_params["label_nc"]
-            else:
-                input_shape_seg[1] = autoencoder_params["label_nc"]
-            input_seg = torch.randn(input_shape_seg).to(device)
-            prediction = inferer(
-                inputs=input,
-                autoencoder_model=stage_1,
-                diffusion_model=stage_2,
-                noise=noise,
-                timesteps=timesteps,
-                seg=input_seg,
-            )
-        else:
-            prediction = inferer(
-                inputs=input, autoencoder_model=stage_1, diffusion_model=stage_2, noise=noise, timesteps=timesteps
+        for scheduler in [DDPMScheduler(num_train_timesteps=10), RFlowScheduler(num_train_timesteps=1000)]:
+            # We infer the VAE shape
+            autoencoder_latent_shape = [i // (2 ** (len(autoencoder_params["channels"]) - 1)) for i in input_shape[2:]]
+            inferer = LatentDiffusionInferer(
+                scheduler=scheduler,
+                scale_factor=1.0,
+                ldm_latent_shape=list(latent_shape[2:]),
+                autoencoder_latent_shape=autoencoder_latent_shape,
             )
-        self.assertEqual(prediction.shape, latent_shape)
+            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():
+                    input_shape_seg[1] = stage_2_params["label_nc"]
+                else:
+                    input_shape_seg[1] = autoencoder_params["label_nc"]
+                input_seg = torch.randn(input_shape_seg).to(device)
+                prediction = inferer(
+                    inputs=input,
+                    autoencoder_model=stage_1,
+                    diffusion_model=stage_2,
+                    noise=noise,
+                    timesteps=timesteps,
+                    seg=input_seg,
+                )
+            else:
+                prediction = inferer(
+                    inputs=input, autoencoder_model=stage_1, diffusion_model=stage_2, noise=noise, timesteps=timesteps
+                )
+            self.assertEqual(prediction.shape, latent_shape)
 
     @parameterized.expand(TEST_CASES_DIFF_SHAPES)
     @skipUnless(has_einops, "Requires einops")
@@ -797,40 +800,42 @@ class TestDiffusionSamplingInferer(unittest.TestCase):
         stage_2.eval()
 
         noise = torch.randn(latent_shape).to(device)
-        scheduler = DDPMScheduler(num_train_timesteps=10)
-        # We infer the VAE shape
-        if ae_model_type == "VQVAE":
-            autoencoder_latent_shape = [i // (2 ** (len(autoencoder_params["channels"]))) for i in input_shape[2:]]
-        else:
-            autoencoder_latent_shape = [i // (2 ** (len(autoencoder_params["channels"]) - 1)) for i in input_shape[2:]]
-
-        inferer = LatentDiffusionInferer(
-            scheduler=scheduler,
-            scale_factor=1.0,
-            ldm_latent_shape=list(latent_shape[2:]),
-            autoencoder_latent_shape=autoencoder_latent_shape,
-        )
-        scheduler.set_timesteps(num_inference_steps=10)
-
-        if dm_model_type == "SPADEDiffusionModelUNet" or ae_model_type == "SPADEAutoencoderKL":
-            input_shape_seg = list(input_shape)
-            if "label_nc" in stage_2_params.keys():
-                input_shape_seg[1] = stage_2_params["label_nc"]
+        for scheduler in [DDPMScheduler(num_train_timesteps=10), RFlowScheduler(num_train_timesteps=1000)]:
+            # We infer the VAE shape
+            if ae_model_type == "VQVAE":
+                autoencoder_latent_shape = [i // (2 ** (len(autoencoder_params["channels"]))) for i in input_shape[2:]]
             else:
-                input_shape_seg[1] = autoencoder_params["label_nc"]
-            input_seg = torch.randn(input_shape_seg).to(device)
-            prediction, _ = inferer.sample(
-                autoencoder_model=stage_1,
-                diffusion_model=stage_2,
-                input_noise=noise,
-                save_intermediates=True,
-                seg=input_seg,
-            )
-        else:
-            prediction = inferer.sample(
-                autoencoder_model=stage_1, diffusion_model=stage_2, input_noise=noise, save_intermediates=False
+                autoencoder_latent_shape = [
+                    i // (2 ** (len(autoencoder_params["channels"]) - 1)) for i in input_shape[2:]
+                ]
+
+            inferer = LatentDiffusionInferer(
+                scheduler=scheduler,
+                scale_factor=1.0,
+                ldm_latent_shape=list(latent_shape[2:]),
+                autoencoder_latent_shape=autoencoder_latent_shape,
             )
-        self.assertEqual(prediction.shape, input_shape)
+            scheduler.set_timesteps(num_inference_steps=10)
+
+            if dm_model_type == "SPADEDiffusionModelUNet" or ae_model_type == "SPADEAutoencoderKL":
+                input_shape_seg = list(input_shape)
+                if "label_nc" in stage_2_params.keys():
+                    input_shape_seg[1] = stage_2_params["label_nc"]
+                else:
+                    input_shape_seg[1] = autoencoder_params["label_nc"]
+                input_seg = torch.randn(input_shape_seg).to(device)
+                prediction, _ = inferer.sample(
+                    autoencoder_model=stage_1,
+                    diffusion_model=stage_2,
+                    input_noise=noise,
+                    save_intermediates=True,
+                    seg=input_seg,
+                )
+            else:
+                prediction = inferer.sample(
+                    autoencoder_model=stage_1, diffusion_model=stage_2, input_noise=noise, save_intermediates=False
+                )
+            self.assertEqual(prediction.shape, input_shape)
 
     @skipUnless(has_einops, "Requires einops")
     def test_incompatible_spade_setup(self):
@@ -866,18 +871,19 @@ class TestDiffusionSamplingInferer(unittest.TestCase):
         stage_2.eval()
         noise = torch.randn((1, 3, 4, 4)).to(device)
         input_seg = torch.randn((1, 3, 8, 8)).to(device)
-        scheduler = DDPMScheduler(num_train_timesteps=10)
-        inferer = LatentDiffusionInferer(scheduler=scheduler, scale_factor=1.0)
-        scheduler.set_timesteps(num_inference_steps=10)
 
-        with self.assertRaises(ValueError):
-            _ = inferer.sample(
-                input_noise=noise,
-                autoencoder_model=stage_1,
-                diffusion_model=stage_2,
-                scheduler=scheduler,
-                seg=input_seg,
-            )
+        for scheduler in [DDPMScheduler(num_train_timesteps=10), RFlowScheduler(num_train_timesteps=1000)]:
+            inferer = LatentDiffusionInferer(scheduler=scheduler, scale_factor=1.0)
+            scheduler.set_timesteps(num_inference_steps=10)
+
+            with self.assertRaises(ValueError):
+                _ = inferer.sample(
+                    input_noise=noise,
+                    autoencoder_model=stage_1,
+                    diffusion_model=stage_2,
+                    scheduler=scheduler,
+                    seg=input_seg,
+                )
 
 
 if __name__ == "__main__":

tests/networks/schedulers/test_scheduler_rflow.py

@@ -0,0 +1,105 @@
+# Copyright (c) MONAI Consortium
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#     http://www.apache.org/licenses/LICENSE-2.0
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from __future__ import annotations
+
+import unittest
+
+import torch
+from parameterized import parameterized
+
+from monai.networks.schedulers import RFlowScheduler
+from tests.test_utils import assert_allclose
+
+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)]
+    )
+
+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_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)]
+    )
+
+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_CASES = TEST_2D_CASE + TEST_3D_CASE
+
+TEST_FULl_LOOP = [
+    [{"sample_method": "uniform"}, (1, 1, 2, 2), torch.Tensor([[[[-0.786166, -0.057519], [2.442662, -0.407664]]]])]
+]
+
+
+class TestRFlowScheduler(unittest.TestCase):
+    @parameterized.expand(TEST_CASES)
+    def test_add_noise(self, input_param, input_shape, expected_shape):
+        scheduler = RFlowScheduler(**input_param)
+        original_sample = torch.zeros(input_shape)
+        timesteps = scheduler.sample_timesteps(original_sample)
+        noise = torch.randn_like(original_sample)
+        timesteps = torch.randint(0, scheduler.num_train_timesteps, (original_sample.shape[0],)).long()
+        noisy = scheduler.add_noise(original_samples=original_sample, noise=noise, timesteps=timesteps)
+        self.assertEqual(noisy.shape, expected_shape)
+
+    @parameterized.expand(TEST_CASES)
+    def test_step_shape(self, input_param, input_shape, expected_shape):
+        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, ...]))
+        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)
+        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, ...]))
+        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)
+        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)
+
+
+if __name__ == "__main__":
+    unittest.main()