[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,