diff --git a/monai/networks/schedulers/rectified_flow.py b/monai/networks/schedulers/rectified_flow.py index f090ebc1e80..8019d087509 100644 --- a/monai/networks/schedulers/rectified_flow.py +++ b/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, ) diff --git a/tests/networks/schedulers/test_scheduler_rflow.py b/tests/networks/schedulers/test_scheduler_rflow.py index 08f4ed3730d..01df788648a 100644 --- a/tests/networks/schedulers/test_scheduler_rflow.py +++ b/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):