From 554ceeac4c900ce946df4e157ad4a84e1ca73b39 Mon Sep 17 00:00:00 2001 From: Don Yin Date: Tue, 20 Oct 2026 19:41:46 +0100 Subject: [PATCH 1/2] Honor transform_scale in RFlowScheduler timestep transforms Pass the configured transformation scale to inference and training timestep transforms instead of always using the helper's default. Add deterministic regression coverage for both paths while preserving the default scale. Signed-off-by: Don Yin Assisted-by: Claude --- monai/networks/schedulers/rectified_flow.py | 2 + .../schedulers/test_scheduler_rflow.py | 48 +++++++++++++++++++ 2 files changed, 50 insertions(+) diff --git a/monai/networks/schedulers/rectified_flow.py b/monai/networks/schedulers/rectified_flow.py index f090ebc1e80..ecc7bf27897 100644 --- a/monai/networks/schedulers/rectified_flow.py +++ b/monai/networks/schedulers/rectified_flow.py @@ -239,6 +239,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 +275,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..1841fe33094 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,47 @@ 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_set_timesteps_with_num_inference_steps_bigger_than_num_train_timesteps(self): scheduler = RFlowScheduler(num_train_timesteps=1000) with self.assertRaises(ValueError): From b3281b18b18eb1a4f5e1e180f13a59388131d367 Mon Sep 17 00:00:00 2001 From: Don Yin Date: Thu, 22 Oct 2026 01:07:06 +0100 Subject: [PATCH 2/2] Reject non-positive transform_scale when timestep transform is on A zero or negative scale makes timestep_transform divide by zero on the first inference timestep. Raise a ValueError in the constructor instead; zero is still accepted when the transform is disabled, since the value is unused then. Signed-off-by: Don Yin Assisted-by: Claude --- monai/networks/schedulers/rectified_flow.py | 2 ++ tests/networks/schedulers/test_scheduler_rflow.py | 7 +++++++ 2 files changed, 9 insertions(+) diff --git a/monai/networks/schedulers/rectified_flow.py b/monai/networks/schedulers/rectified_flow.py index ecc7bf27897..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 diff --git a/tests/networks/schedulers/test_scheduler_rflow.py b/tests/networks/schedulers/test_scheduler_rflow.py index 1841fe33094..01df788648a 100644 --- a/tests/networks/schedulers/test_scheduler_rflow.py +++ b/tests/networks/schedulers/test_scheduler_rflow.py @@ -143,6 +143,13 @@ def test_sample_timesteps_transform_scale(self, transform_scale, spatial_dim, us ) 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):