[Bugfix][BAGEL] Restore original denoise step count (#7049)

Signed-off-by: Trigger <Meng.Bo.Wang@outlook.com>
Signed-off-by: Gao Han <hgaoaf@connect.ust.hk>
Co-authored-by: Gao Han <hgaoaf@connect.ust.hk>

tests/diffusion/models/bagel/test_step_execution.py

@@ -41,7 +41,7 @@ class _Sampling:
 
 @dataclass
 class _ScheduleConfig:
-    _denoise_schedule_extra_step: bool = True
+    _denoise_schedule_extra_step: bool = False
 
 
 @dataclass
@@ -245,8 +245,8 @@ def test_prepare_encode_populates_request_local_state_and_schedule():
     assert result is state
     assert state.latents is prepared["generation_input"]["packed_init_noises"]
     assert state.latents.dtype == torch.float32
-    assert state.total_steps == 5
-    expected_points = torch.linspace(1, 0, 6)
+    assert state.total_steps == 4
+    expected_points = torch.linspace(1, 0, 5)
     expected_points = 2.0 * expected_points / (1 + expected_points)
     assert torch.equal(state.timesteps, expected_points[:-1])
     assert state.timesteps.dtype == torch.float32
@@ -264,27 +264,28 @@ def test_prepare_encode_populates_request_local_state_and_schedule():
     assert "packed_init_noises" not in state.extra["bagel_generation_input"]
 
 
-def test_prepare_encode_rejects_non_positive_step_schedule():
-    prepared = _prepared_context()
-    prepared["gen_params"].num_timesteps = 0
-    pipeline = _pipeline_for_prepare(prepared)
-    state = StepRequestState(request_id="req", sampling=_sampling(), prompt="draw a cat")
-
-    with pytest.raises(ValueError, match="num_inference_steps >= 1"):
-        pipeline.prepare_encode(state)
-
-
-def test_prepare_encode_supports_one_denoising_update():
+def test_prepare_encode_rejects_one_step_schedule():
     prepared = _prepared_context()
     prepared["gen_params"].num_timesteps = 1
     pipeline = _pipeline_for_prepare(prepared)
     state = StepRequestState(request_id="req", sampling=_sampling(), prompt="draw a cat")
 
-    pipeline.prepare_encode(state)
+    with pytest.raises(ValueError, match="num_inference_steps >= 2"):
+        pipeline.prepare_encode(state)
 
-    assert state.total_steps == 1
-    assert len(state.timesteps) == 1
-    assert len(state.extra["bagel_dts"]) == 1
+
+def test_lance_style_extra_schedule_point_supports_one_denoising_update():
+    latents = torch.zeros(1, 3)
+
+    timesteps, dts = Bagel.prepare_denoise_schedule(
+        _ScheduleConfig(_denoise_schedule_extra_step=True),
+        latents,
+        num_timesteps=1,
+        timestep_shift=1.0,
+    )
+
+    assert len(timesteps) == 1
+    assert len(dts) == 1
 
 
 def test_prepare_encode_rejects_sequence_parallel_before_forward():
@@ -520,7 +521,7 @@ def test_complete_step_updates_match_full_bagel_loop(cfg_text_scale: float, velo
 
     bagel = MagicMock()
     bagel._sp_size = 1
-    bagel._denoise_schedule_extra_step = True
+    bagel._denoise_schedule_extra_step = False
     bagel.prepare_denoise_schedule = types.MethodType(Bagel.prepare_denoise_schedule, bagel)
     bagel.forward.side_effect = lambda x_t, **_kwargs: torch.full(
         x_t.shape,

tests/diffusion/models/bagel/test_trajectory_recording.py

@@ -19,14 +19,18 @@ pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 NUM_TOKENS = 8
 HIDDEN_DIM = 16
 NUM_TIMESTEPS = 5
-EXPECTED_STEPS = NUM_TIMESTEPS
+# Official BAGEL samples num_timesteps points then drops the terminal t=0,
+# yielding num_timesteps - 1 denoise steps.
+EXPECTED_STEPS = NUM_TIMESTEPS - 1
 
 
 def _make_mock_bagel(mocker: MockerFixture):
     """Create a mock Bagel with forward returning constant velocity."""
     mock = mocker.MagicMock(spec=Bagel)
     mock._sp_size = 1
-    mock._denoise_schedule_extra_step = True
+    # MagicMock would otherwise auto-return a truthy stub for this flag; pin it
+    # to the official BAGEL schedule convention (num_timesteps points).
+    mock._denoise_schedule_extra_step = False
 
     # forward returns a small constant velocity so x_t changes each step
     def fake_forward(self, x_t, **kwargs):
@@ -93,6 +97,14 @@ def bagel_and_args(
 class TestTrajectoryRecording:
     """Tests for trajectory latent/timestep recording in generate_image."""
 
+    def test_complete_request_rejects_one_step_schedule(self, bagel_and_args):
+        """The non-step image path must not decode untouched initial noise."""
+        bagel, args = bagel_and_args
+        args["num_timesteps"] = 1
+
+        with pytest.raises(ValueError, match="num_inference_steps >= 2"):
+            bagel.generate_image(**args)
+
     def test_trajectory_disabled_returns_none(self, bagel_and_args):
         bagel, args = bagel_and_args
 
@@ -178,7 +190,7 @@ class TestTrajectoryRecording:
         bagel, args = bagel_and_args
 
         # Recompute the expected timestep schedule (mirrors generate_image logic)
-        ts = torch.linspace(1, 0, args["num_timesteps"] + 1)
+        ts = torch.linspace(1, 0, args["num_timesteps"])
         shift = args.get("timestep_shift", 1.0)
         ts = shift * ts / (1 + (shift - 1) * ts)
         expected_timesteps = ts[:-1]  # last element is dropped
@@ -196,7 +208,7 @@ class TestTrajectoryRecording:
         bagel, args = bagel_and_args
 
         # Compute the number of denoising steps the same way generate_image does
-        ts = torch.linspace(1, 0, args["num_timesteps"] + 1)
+        ts = torch.linspace(1, 0, args["num_timesteps"])
         shift = args.get("timestep_shift", 1.0)
         ts = shift * ts / (1 + (shift - 1) * ts)
         num_steps = len(ts) - 1  # timesteps = ts[:-1]