[BugFix]: Fix Bagel generate_image denoising schedule (#4509)

tests/diffusion/models/bagel/test_trajectory_recording.py

@@ -19,15 +19,18 @@ pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 NUM_TOKENS = 8
 HIDDEN_DIM = 16
 NUM_TIMESTEPS = 5
-# generate_image samples (num_timesteps + 1) points then drops the last, matching
-# upstream lance: see Bagel.generate_image in bagel_transformer.py.
-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
+    # 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):
@@ -179,7 +182,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
@@ -197,7 +200,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]

tests/distributed/omni_connectors/test_bagel_mooncake_connector.py

@@ -32,7 +32,7 @@ BAGEL_MOONCAKE_CI_DEPLOY = get_deploy_config_path("ci/bagel_mooncake.yaml")
 
 # Reference pixel data extracted from the known-good output image
 # Each entry contains (x, y) position and expected (R, G, B) values
-# "Generated with seed=52, num_inference_steps=14,
+# "Generated with seed=52, num_inference_steps=15,
 # prompt='A cute cat'"
 REFERENCE_PIXELS = [
     {"position": (100, 100), "rgb": (115, 113, 94)},
@@ -63,7 +63,7 @@ def _find_free_port() -> int:
     return port
 
 
-def _configure_sampling_params(omni: Omni, num_inference_steps: int = 14) -> list:
+def _configure_sampling_params(omni: Omni, num_inference_steps: int = 15) -> list:
     """Configure sampling parameters for Bagel text2img generation.
 
     Args:

tests/distributed/omni_connectors/test_bagel_shared_memory_connector.py

@@ -30,7 +30,7 @@ pytestmark = [pytest.mark.usefixtures("clean_gpu_memory_between_tests")]
 BAGEL_CI_DEPLOY = get_deploy_config_path("ci/bagel.yaml")
 
 # Reference pixel data extracted from the known-good output image
-# Generated with seed=52, num_inference_steps=14,
+# Generated with seed=52, num_inference_steps=15,
 # prompt='Change the grass color to red',
 # input image: 2560px-Gfp-wisconsin-madison-the-nature-boardwalk.jpg
 REFERENCE_PIXELS = [
@@ -48,7 +48,7 @@ REFERENCE_PIXELS = [
 
 
 # text2img reference pixels (aligned with offline `bagel/end2end.py` text2img, 15 steps)
-# "Generated with seed=52, num_inference_steps=14,
+# "Generated with seed=52, num_inference_steps=15,
 # prompt='A cute cat'"
 TEXT2IMG_REFERENCE_PIXELS = [
     {"position": (100, 100), "rgb": (115, 113, 94)},
@@ -88,7 +88,7 @@ def _find_free_port() -> int:
     return port
 
 
-def _configure_sampling_params(omni: Omni, num_inference_steps: int = 14) -> list:
+def _configure_sampling_params(omni: Omni, num_inference_steps: int = 15) -> list:
     """Configure sampling parameters for Bagel img2img generation.
 
     Args: