Honor transform_scale in RFlowScheduler timestep transforms

monai/networks/schedulers/rectified_flow.py

@@ -174,6 +174,8 @@ def __init__(
 
         # timestep transform
         self.use_timestep_transform = use_timestep_transform
+        if use_timestep_transform and transform_scale <= 0:
+            raise ValueError("`transform_scale` must be positive when `use_timestep_transform` is True.")
         self.transform_scale = transform_scale
         self.steps_offset = steps_offset
 
@@ -239,6 +241,7 @@ def set_timesteps(
                     t,
                     input_img_size_numel=input_img_size_numel,
                     base_img_size_numel=self.base_img_size_numel,
+                    scale=self.transform_scale,
                     num_train_timesteps=self.num_train_timesteps,
                     spatial_dim=self.spatial_dim,
                 )
@@ -274,6 +277,7 @@ def sample_timesteps(self, x_start):
                 t,
                 input_img_size_numel=input_img_size_numel,
                 base_img_size_numel=self.base_img_size_numel,
+                scale=self.transform_scale,
                 num_train_timesteps=self.num_train_timesteps,
                 spatial_dim=len(x_start.shape) - 2,
             )

tests/networks/schedulers/test_scheduler_rflow.py

@@ -52,6 +52,13 @@
 
 TEST_CASES = TEST_2D_CASE + TEST_3D_CASE
 
+TEST_TRANSFORM_SCALE_CASES = [
+    (transform_scale, spatial_dim, use_discrete_timesteps)
+    for transform_scale in [0.5, 1.0, 2.0]
+    for spatial_dim in [2, 3]
+    for use_discrete_timesteps in [False, True]
+]
+
 TEST_FULl_LOOP = [
     [{"sample_method": "uniform"}, (1, 1, 2, 2), torch.Tensor([[[[-0.786166, -0.057519], [2.442662, -0.407664]]]])]
 ]
@@ -95,6 +102,54 @@ def test_set_timesteps(self):
         self.assertEqual(scheduler.num_inference_steps, 100)
         self.assertEqual(len(scheduler.timesteps), 100)
 
+    @parameterized.expand(TEST_TRANSFORM_SCALE_CASES)
+    def test_set_timesteps_transform_scale(self, transform_scale, spatial_dim, use_discrete_timesteps):
+        scheduler = RFlowScheduler(
+            use_timestep_transform=True,
+            transform_scale=transform_scale,
+            use_discrete_timesteps=use_discrete_timesteps,
+            base_img_size_numel=4**spatial_dim,
+            spatial_dim=spatial_dim,
+        )
+        scheduler.set_timesteps(num_inference_steps=4, input_img_size_numel=4**spatial_dim)
+        original_timesteps = torch.tensor([1.0, 0.75, 0.5, 0.25])
+        expected = 1000 * transform_scale * original_timesteps / (1 + (transform_scale - 1) * original_timesteps)
+        expected = expected.to(torch.float16)
+        if use_discrete_timesteps:
+            expected = expected.long()
+        assert_allclose(scheduler.timesteps, expected, rtol=0, atol=0)
+
+    @parameterized.expand(TEST_TRANSFORM_SCALE_CASES)
+    def test_sample_timesteps_transform_scale(self, transform_scale, spatial_dim, use_discrete_timesteps):
+        inputs = torch.zeros((8, 1) + (4,) * spatial_dim)
+        for sample_method in ["uniform", "logit-normal"]:
+            with self.subTest(sample_method=sample_method):
+                params = {
+                    "sample_method": sample_method,
+                    "transform_scale": transform_scale,
+                    "use_discrete_timesteps": use_discrete_timesteps,
+                    "base_img_size_numel": 4**spatial_dim,
+                    "spatial_dim": spatial_dim,
+                }
+                original_scheduler = RFlowScheduler(**params)
+                scheduler = RFlowScheduler(**params, use_timestep_transform=True)
+                with torch.random.fork_rng(devices=[]):
+                    torch.manual_seed(42)
+                    original_timesteps = original_scheduler.sample_timesteps(inputs) / 1000
+                    torch.manual_seed(42)
+                    timesteps = scheduler.sample_timesteps(inputs)
+                expected = (
+                    1000 * transform_scale * original_timesteps / (1 + (transform_scale - 1) * original_timesteps)
+                )
+                assert_allclose(timesteps, expected)
+
+    def test_non_positive_transform_scale(self):
+        for transform_scale in (0.0, -1.0):
+            with self.assertRaises(ValueError):
+                RFlowScheduler(use_timestep_transform=True, transform_scale=transform_scale)
+        scheduler = RFlowScheduler(use_timestep_transform=False, transform_scale=0.0)
+        self.assertEqual(scheduler.transform_scale, 0.0)
+
     def test_set_timesteps_with_num_inference_steps_bigger_than_num_train_timesteps(self):
         scheduler = RFlowScheduler(num_train_timesteps=1000)
         with self.assertRaises(ValueError):