diff --git a/tests/diffusion/models/bagel/test_trajectory_recording.py b/tests/diffusion/models/bagel/test_trajectory_recording.py index aa86f033a..f3dbc0df2 100644 --- a/tests/diffusion/models/bagel/test_trajectory_recording.py +++ b/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] diff --git a/tests/distributed/omni_connectors/test_bagel_mooncake_connector.py b/tests/distributed/omni_connectors/test_bagel_mooncake_connector.py index 106507114..cac7d7ed0 100644 --- a/tests/distributed/omni_connectors/test_bagel_mooncake_connector.py +++ b/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: diff --git a/tests/distributed/omni_connectors/test_bagel_shared_memory_connector.py b/tests/distributed/omni_connectors/test_bagel_shared_memory_connector.py index eb5a79d6a..19457f0f2 100644 --- a/tests/distributed/omni_connectors/test_bagel_shared_memory_connector.py +++ b/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: diff --git a/vllm_omni/diffusion/models/bagel/bagel_transformer.py b/vllm_omni/diffusion/models/bagel/bagel_transformer.py index ba3a692bc..bc3eef49d 100644 --- a/vllm_omni/diffusion/models/bagel/bagel_transformer.py +++ b/vllm_omni/diffusion/models/bagel/bagel_transformer.py @@ -1129,6 +1129,13 @@ class Bagel(nn.Module): config_class = BagelConfig base_model_prefix = "bagel" + # Flow-matching denoise schedule convention. Official BAGEL samples + # ``num_timesteps`` points over [1, 0] and drops the terminal t=0, yielding + # ``num_timesteps - 1`` Euler steps. Lance samples one extra point + # (``num_timesteps + 1``) for ``num_timesteps`` steps; ``LanceBagel`` flips + # this on. See https://github.com/vllm-project/vllm-omni/issues/4470. + _denoise_schedule_extra_step: bool = False + def __init__( self, language_model, @@ -1690,11 +1697,11 @@ class Bagel(nn.Module): frame_condition_token_indexes = frame_condition_token_indexes.to(x_t.device).long() pinned_x_t = x_t[frame_condition_token_indexes].clone() - # Use num_timesteps + 1 sample points so we get `num_timesteps` denoise - # steps after dropping the terminal t=0 (which has no dt). Upstream - # Lance / BAGEL both use this convention; without the +1 we silently - # run one fewer denoise iteration than the user asked for. - timesteps = torch.linspace(1, 0, num_timesteps + 1, device=x_t.device) + # Build the flow-matching schedule. BAGEL drops the terminal t=0 for + # ``num_timesteps - 1`` Euler steps; Lance keeps it for ``num_timesteps``. + # ``_denoise_schedule_extra_step`` (overridden by ``LanceBagel``) selects which. + num_sample_points = num_timesteps + 1 if self._denoise_schedule_extra_step else num_timesteps + timesteps = torch.linspace(1, 0, num_sample_points, device=x_t.device) timesteps = timestep_shift * timesteps / (1 + (timestep_shift - 1) * timesteps) dts = timesteps[:-1] - timesteps[1:] timesteps = timesteps[:-1] diff --git a/vllm_omni/diffusion/models/lance/lance_transformer.py b/vllm_omni/diffusion/models/lance/lance_transformer.py index 012ba3f42..12959eb64 100644 --- a/vllm_omni/diffusion/models/lance/lance_transformer.py +++ b/vllm_omni/diffusion/models/lance/lance_transformer.py @@ -278,6 +278,11 @@ class LanceBagel(Bagel): ``forward_cache_update_vae`` — is reused unchanged. """ + # Upstream Lance samples ``num_timesteps + 1`` schedule points (one extra + # over official BAGEL) so the denoise loop runs exactly ``num_timesteps`` + # Euler steps. See Bagel.generate_image and issue #4470. + _denoise_schedule_extra_step: bool = True + @staticmethod def _lance_compute_vit_buckets(): """Replicate upstream Lance's ``BucketResize.init_buckets`` for the ViT path.