[Misc] Clean up method name in BAGEL. (#2501)

Signed-off-by: Zhengyuan Su <su.zhengyuan@u.nus.edu>
Co-authored-by: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

tests/diffusion/models/bagel/test_trajectory_recording.py

@@ -24,15 +24,15 @@ EXPECTED_STEPS = NUM_TIMESTEPS - 1
 
 
 def _make_mock_bagel():
-    """Create a mock Bagel with _forward_flow returning constant velocity."""
+    """Create a mock Bagel with forward returning constant velocity."""
     mock = MagicMock(spec=Bagel)
     mock._sp_size = 1
 
-    # _forward_flow returns a small constant velocity so x_t changes each step
-    def fake_forward_flow(self, x_t, **kwargs):
+    # forward returns a small constant velocity so x_t changes each step
+    def fake_forward(self, x_t, **kwargs):
         return torch.ones_like(x_t) * 0.1
 
-    mock._forward_flow = types.MethodType(fake_forward_flow, mock)
+    mock.forward = types.MethodType(fake_forward, mock)
     # _merge_naive_caches is called in the batched CFG path
     mock._merge_naive_caches = types.MethodType(lambda self, caches: NaiveCache(1), mock)
 

vllm_omni/diffusion/cache/teacache/backend.py

@@ -48,16 +48,7 @@ def enable_bagel_teacache(pipeline: Any, config: DiffusionCacheConfig) -> None:
         coefficients=config.coefficients,
     )
     transformer = pipeline.bagel
-    original_forward_flow = transformer._forward_flow
-
-    import types
-
-    def forward_alias(self, *args, **kwargs):
-        return original_forward_flow(*args, **kwargs)
-
-    transformer.forward = types.MethodType(forward_alias, transformer)
     apply_teacache_hook(transformer, teacache_config)
-    transformer._forward_flow = transformer.forward
     pipeline.transformer = transformer
 
     logger.info(

vllm_omni/diffusion/cache/teacache/coefficient_estimator.py

@@ -1,7 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
-import types
 from typing import Any
 
 import numpy as np
@@ -74,15 +73,8 @@ class BagelAdapter:
 
     @staticmethod
     def install_hook(transformer: Any, hook: DataCollectionHook) -> None:
-        original_forward_flow = transformer._forward_flow
-
-        def forward_alias(self, *args, **kwargs):
-            return original_forward_flow(*args, **kwargs)
-
-        transformer.forward = types.MethodType(forward_alias, transformer)
         registry = HookRegistry.get_or_create(transformer)
         registry.register_hook(hook._HOOK_NAME, hook)
-        transformer._forward_flow = transformer.forward
 
 
 class StableAudioAdapter:

vllm_omni/diffusion/models/bagel/bagel_transformer.py

@@ -1734,7 +1734,7 @@ class Bagel(nn.Module):
                     packed_seqlens=packed_seqlens,
                 )
 
-                v_t = self._forward_flow_single_branch(
+                v_t = self.forward_single_branch(
                     **common,
                     packed_indexes=packed_indexes,
                     packed_position_ids=packed_position_ids,
@@ -1744,7 +1744,7 @@ class Bagel(nn.Module):
                 )
 
                 if cfg_text_scale_ > 1.0:
-                    cfg_text_v_t = self._forward_flow_single_branch(
+                    cfg_text_v_t = self.forward_single_branch(
                         **common,
                         packed_indexes=cfg_text_packed_query_indexes,
                         packed_position_ids=cfg_text_packed_position_ids,
@@ -1754,7 +1754,7 @@ class Bagel(nn.Module):
                     )
                     cfg_img_v_t = None
                     if cfg_img_scale_ > 1.0:
-                        cfg_img_v_t = self._forward_flow_single_branch(
+                        cfg_img_v_t = self.forward_single_branch(
                             **common,
                             packed_indexes=cfg_img_packed_query_indexes,
                             packed_position_ids=cfg_img_packed_position_ids,
@@ -1790,7 +1790,7 @@ class Bagel(nn.Module):
         if use_sp:
             for i, t in enumerate(timesteps):
                 timestep = torch.tensor([t] * x_t.shape[0], device=x_t.device)
-                v_t = self._forward_flow_single_branch(
+                v_t = self.forward_single_branch(
                     x_t=x_t,
                     timestep=timestep,
                     packed_vae_token_indexes=packed_vae_token_indexes,
@@ -1883,7 +1883,7 @@ class Bagel(nn.Module):
             else:
                 cfg_text_scale_ = 1.0
                 cfg_img_scale_ = 1.0
-            v_t = self._forward_flow(
+            v_t = self.forward(
                 x_t=x_t,
                 timestep=timestep,
                 packed_vae_token_indexes=packed_vae_token_indexes,
@@ -2019,7 +2019,7 @@ class Bagel(nn.Module):
 
             if use_cfg_this_step:
                 # CFG interval: each rank computes its own branch
-                local_v_t = self._forward_flow_single_branch(
+                local_v_t = self.forward_single_branch(
                     x_t=x_t,
                     timestep=timestep,
                     packed_vae_token_indexes=packed_vae_token_indexes,
@@ -2046,7 +2046,7 @@ class Bagel(nn.Module):
                 )
             else:
                 # Outside CFG interval: all ranks compute with gen inputs, no comm
-                v_t = self._forward_flow_single_branch(
+                v_t = self.forward_single_branch(
                     x_t=x_t,
                     timestep=timestep,
                     packed_vae_token_indexes=packed_vae_token_indexes,
@@ -2128,7 +2128,7 @@ class Bagel(nn.Module):
 
         return v_t
 
-    def _forward_flow_single_branch(
+    def forward_single_branch(
         self,
         x_t: torch.Tensor,
         timestep: torch.LongTensor,
@@ -2258,7 +2258,7 @@ class Bagel(nn.Module):
         v_t = v_t[packed_vae_token_indexes]
         return v_t
 
-    def _forward_flow(
+    def forward(
         self,
         x_t: torch.Tensor,
         timestep: torch.LongTensor,

vllm_omni/diffusion/models/bagel/pipeline_bagel.py

@@ -159,8 +159,8 @@ class BagelPipeline(nn.Module, DiffusionPipelineProfilerMixin):
         self.od_config = od_config
         self.device = get_local_device()
 
-        self._scheduler: object | None = None
-        self._scheduler_kwargs: dict = {}
+        self.scheduler: object | None = None
+        self.scheduler_kwargs: dict = {}
 
         model = od_config.model
         local_files_only = os.path.exists(model)
@@ -654,8 +654,8 @@ class BagelPipeline(nn.Module, DiffusionPipelineProfilerMixin):
                 cfg_img_key_values_lens=generation_input_cfg_img["cfg_key_values_lens"],
                 cfg_img_packed_key_value_indexes=generation_input_cfg_img["cfg_packed_key_value_indexes"],
                 return_trajectory_latents=req.sampling_params.return_trajectory_latents,
-                scheduler=self._scheduler,
-                scheduler_kwargs=self._scheduler_kwargs,
+                scheduler=self.scheduler,
+                scheduler_kwargs=self.scheduler_kwargs,
             )
 
         img = self._decode_image_from_latent(self.bagel, self.vae, latents[0], image_shape)